Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean
  Field Neural Networks
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 $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.
Blake Bordelon, Cengiz Pehlevan
# Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks

###### 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.

## 1 Introduction

Learning dynamics of deep neural networks are challenging to analyze and understand theoretically, but recent progress has been made by studying the idealization of infinite-width networks. Two types of infinite-width limits have been especially fruitful. First, the kernel or lazy infinite-width limit, which arises in the standard or neural tangent kernel (NTK) parameterization, gives prediction dynamics which correspond to a linear model [1, 2, 3, 4, 5]. This limit is theoretically tractable but fails 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 \(\mu\)-parameterization allows feature learning at infinite width [6, 7, 8, 9].

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 is still small in this limit) but also introduces variance in the kernels and the predictions over random initializations [10, 11, 12, 13, 14, 15]. Because of these competing effects, in some situations wider networks are better, and in others wider networks perform worse [16].

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 a methodology for reasoning about detrimental finite-size effects in such feature-learning neural networks. We show that observable averages involving kernels and predictions obey a
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 \(\mathcal{O}(1/\text{width})\) 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:

1. Starting from a dynamical mean field theory (DMFT) description of infinite-width nonlinear deep neural network training dynamics, we provide a 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 \(\mathcal{O}(1/\text{width})\) 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.
3. 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-noise ratio (SNR) of kernels 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 finite-width effects and how richer training can reduce this effect.
5. We show that large learning rate effects such as edge-of-stability [19, 20, 21] dynamics 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.

### Related Works

Infinite-width networks at initialization converge to a Gaussian process with a covariance kernel that is computed with a layerwise recursion [23, 24, 25, 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, 28, 29, 30, 14]. During gradient-based training with the NTK parameterization, a hierarchy of differential equations have been utilized to compute small feature learning corrections to the kernel through training [10, 11, 12, 13]. 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 kernels at finite width during inference [31, 32]. Other relevant work in this domain are [33, 34, 35, 36, 37, 38, 39].

An alternative to standard/NTK parameterization is the mean field or \(\mu P\)-limit where features evolve even at infinite width [6, 7, 8, 9, 40, 41, 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 to initialization noise. Prior work has studied how the weight space distribution and predictions converge to mean field dynamics with a dynamical error \(\mathcal{O}(1/\sqrt{\text{width}})\)[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 [9, 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 a perturbation expansion around DMFT is a long celebrated technique to obtain finite size corrections in physics [47, 48, 49, 50] and has been utilized for random untrained recurrent networks [51, 52], and more recently to calculate variance of feature kernels \(\Phi^{\ell}\) 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 is large. For a multi-layer perceptron (MLP), the network is defined as a map from input \(\mathbf{x}_{\mu}\in\mathbb{R}^{D}\) to hidden activations \(\mathbf{h}_{\mu}^{\ell}\in\mathbb{R}^{N}\) in layers \(\ell\in\{1,...,L\}\) and finally output \(f_{\mu}\)

\[f_{\mu}=\frac{1}{\gamma N}\mathbf{w}^{L}\cdot\phi(\mathbf{h}_{\mu}^{L})\,\quad\mathbf{h}_{\mu}^{\ell+1}= \frac{1}{\sqrt{N}}\mathbf{W}^{\ell}\phi(\mathbf{h}_{\mu}^{\ell})\,\quad\mathbf{h}_{\mu}^{1}=\frac{1}{\sqrt{D}}\mathbf{W}^{0}\mathbf{x}_{ \mu}, \tag{1}\]

where \(\gamma\) is a scale factor that controls feature learning strength, with large \(\gamma\) leading to rich feature learning dynamics and the limit of small \(\gamma\to 0\) (or generally if \(\gamma\) scales as \(N^{-\alpha}\) for \(\alpha>0\) as \(N\to\infty\), NTK parameterization corresponds to \(\alpha=\frac{1}{2}\)) gives lazy learning where no features are learned [4, 7, 9]. The parameters \(\mathbf{\theta}=\{\mathbf{W}^{0},\mathbf{W}^{1},...,\mathbf{w}^{L}\}\) are optimized with gradient descent or gradient flow \(\frac{d}{dt}\mathbf{\theta}=-N\gamma^{2}\nabla_{\mathbf{\theta}}\mathcal{L}\) where \(\mathcal{L}=\mathbb{E}_{\mathbf{x}_{\mu}\in\mathcal{D}}\left({f(\mathbf{x}_{\mu},\mathbf{ \theta}),y_{\mu}}\right)\) is a loss computed over dataset \(\mathcal{D}=\{(\mathbf{x}_{1},y_{1}),(\mathbf{x}_{2},y_{2}),\ldots(\mathbf{x}_{P},y_{P})\}\). This parameterization and learning rate scaling ensures that \(\frac{d}{dt}f_{\mu}\sim\mathcal{O}_{N,\gamma}(1)\) and \(\frac{d}{dt}\mathbf{h}_{\mu}^{\ell}=\mathcal{O}_{N,\gamma}(\gamma)\) at initialization. This is equivalent to maximal update parameterization (\(\mu\)P)[8], which can be easily extended to other architectures including neural networks with trainable bias parameters as well as convolutional, recurrent, and attention layers [8, 9].

## 3 Review of Dynamical Mean Field Theory

The infinite-width training dynamics of feature learning neural networks was described by a DMFT in [9, 46]. We first review the DMFT's key concepts, before extending it to get insight into finite-widths.

To arrive at the DMFT, one first notes that the training dynamics of such networks can be rewritten in terms of a collection of dynamical variables (or _order parameters_) \(\mathbf{q}=\mathrm{Vec}\{f_{\mu}(t),\Phi_{\mu\nu}^{\ell}(t,s),G_{\mu\nu}^{\ell}(t,s ),...\}\)[9], which include feature and gradient kernels [9, 54]

\[\Phi_{\mu\nu}^{\ell}(t,s)=\frac{1}{N}\phi(\mathbf{h}_{\mu}^{\ell}(t))\cdot\phi(\bm {h}_{\nu}^{\ell}(s))\,\quad G_{\mu\nu}^{\ell}(t,s)=\frac{1}{N}\mathbf{g}_{\mu}^{\ell}(t) \cdot\mathbf{g}_{\nu}^{\ell}(s), \tag{2}\]

where \(\mathbf{g}_{\mu}^{\ell}(t)=\gamma N\frac{\partial f_{\mu}(t)}{\partial\mathbf{h}_{\mu }^{\ell}(t)}\) are the back-propagated gradient signals. Further, for width-\(N\) networks the distribution of these dynamical variables across weight initializations (from a Gaussian distribution \(\mathbf{\theta}\sim\mathcal{N}(0,\mathbf{I})\)) is given by \(p(\mathbf{q})\propto\exp{(NS(\mathbf{q}))}\), where the action \(S(\mathbf{q})\) contains interactions between neuron activations and the kernels at each layer [9].

The DMFT introduced in [9] arises in the \(N\to\infty\) limit when \(p(\mathbf{q})\) is strongly peaked around the saddle point \(\mathbf{q}_{\infty}\) where \(\frac{\partial S}{\partial\mathbf{q}}|_{\mathbf{q}_{\infty}}=0\). 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 of neuron pre-activations in each layer (informally the histogram of the elements of \(\mathbf{h}_{\mu}^{\ell}(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_{\mu}^{\ell}(t)\). Kernels in (2) are now computed as _averages_ over these infinite-width single site processes \(\Phi_{\mu\nu}^{\ell}(t,s)=\left\langle\phi(h_{\mu}^{\ell}(t))\phi(h_{\nu}^{ \ell}(s))\right\rangle\), \(G_{\mu\nu}^{\ell}(t,s)=\left\langle g_{\mu}^{\ell}(t)g_{\nu}^{\ell}(s)\right\rangle\), where averages arise from the \(N\to\infty\) limit of the dot products in (2). DMFT also provides a set of self-consistent equations that describe the complete statistics of these random processes, which depend on the kernels, as well as other quantities.

## 4 Dynamical Fluctuations Around Mean Field Theory

We are interested in going beyond the infinite-width limit to study more realistic finite-width networks. In this regime, the order parameters \(\mathbf{q}\) fluctuate in a \(\mathcal{O}(1/\sqrt{N})\) neighborhood of \(\mathbf{q}_{\infty}\)[55, 51, 53, 46]. Statistics of these fluctuations can be calculated from a general cumulant expansion (see App. C) [55, 56, 51]. We will focus on the leading-order corrections to the infinite-width limit in this expansion.

**Proposition 1**_The finite-width \(N\) average of observable \(O(\mathbf{q})\) across initializations, which we denote by \(\left\langle O(\mathbf{q})\right\rangle_{N}\), admits an expansion of the form whose leading terms are_

\[\left\langle O(\mathbf{q})\right\rangle_{N}=\frac{\int d\mathbf{q}\exp{(NS[\mathbf{q}])}O( \mathbf{q})}{\int d\mathbf{q}\exp{(NS[\mathbf{q}])}}=\left\langle O(\mathbf{q})\right\rangle _{\infty}+N\left[\left\langle V(\mathbf{q})O(\mathbf{q})\right\rangle_{\infty}- \left\langle V(\mathbf{q})\right\rangle_{\infty}\left\langle O(\mathbf{q})\right\rangle_ {\infty}\right]+..., \tag{3}\]
_where \(\left\langle\right\rangle_{\infty}\) denotes an average over the Gaussian distribution \(\mathbf{q}\sim\mathcal{N}\left(\mathbf{q}_{\infty},-\frac{1}{N}\left(\nabla_{\mathbf{q}}^{2} S[\mathbf{q}_{\infty}]\right)^{-1}\right)\) and the function \(V(\mathbf{q})\equiv S(\mathbf{q})-S(\mathbf{q}_{\infty})-\frac{1}{2}(\mathbf{q}-\mathbf{q}_{\infty}) ^{\top}\nabla_{\mathbf{q}}^{2}S(\mathbf{q}_{\infty})(\mathbf{q}-\mathbf{q}_{\infty})\) contains cubic and higher terms in the Taylor expansion of \(S\) around \(\mathbf{q}_{\infty}\). The terms shown include all the leading and sub-leading terms in the series in powers of \(1/N\). The terms in ellipses are at least \(\mathcal{O}(N^{-1})\) suppressed compared to the terms provided._

The proof of this statement is given in App. C. The central object to characterize finite size effects is the unperturbed covariance (the _propagator_): \(\mathbf{\Sigma}=-\left[\nabla^{2}S(\mathbf{q}_{\infty})\right]^{-1}\). This object can be shown to capture leading order fluctuation statistics \(\left\langle\left(\mathbf{q}-\mathbf{q}_{\infty}\right)\left(\mathbf{q}-\mathbf{q}_{\infty} \right)^{\top}\right\rangle_{N}=\frac{1}{N}\mathbf{\Sigma}+\mathcal{O}(N^{-2})\) (App. C.1), which can be used to 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 \(\gamma\)[7, 17, 18].

To calculate such corrections, in App. D, we provide a complete description of Hessian \(\nabla_{\mathbf{q}}^{2}S(\mathbf{q})\) and its inverse (the propagator) for a depth-\(L\) network. This description constitutes one of our main results. The resulting expressions are lengthy and are left to App. D. Here, we discuss them at a high level. Conceptually there are two primary ingredients for obtaining the full propagator:

* Hessian sub-blocks \(\kappa\) which describe the _uncoupled variances_ of the kernels, such as \[\kappa_{\mu\nu\alpha\beta}^{\Phi^{\ell}}(t,s,t^{\prime},s^{\prime})\equiv \left\langle\phi(h_{\mu}^{\ell}(t))\phi(h_{\nu}^{\ell}(s))\phi(h_{\alpha}^{ \ell}(t^{\prime}))\phi(h_{\beta}^{\ell}(s^{\prime}))\right\rangle-\Phi_{\mu \nu}^{\ell}(t,s)\Phi_{\alpha\beta}^{\ell}(t^{\prime},s^{\prime})\] (4) Similar terms also appear in other studies on finite width Bayesian inference [13, 31, 32] and in studies on kernel variance at initialization [27, 14, 29, 53].
* Blocks which capture the _sensitivity_ of field averages to pertubations of order parameters, such as \[D_{\mu\nu\alpha\beta}^{\Phi^{\ell}\Phi^{\ell-1}}(t,s,t^{\prime},s^{\prime}) \equiv\frac{\partial\left\langle\phi(h_{\mu}^{\ell}(t))\phi(h_{\nu}^{\ell}(s ))\right\rangle}{\partial\Phi_{\alpha\beta}^{\ell-1}\left(t^{\prime},s^{\prime }\right)}\,\quad D_{\mu\nu\alpha}^{G^{\ell}\Delta}(t,s,t^{\prime})\equiv\frac{ \partial\left\langle g_{\mu}^{\ell}(t)g_{\nu}^{\ell}(s)\right\rangle}{ \partial\Delta_{\alpha}(t^{\prime})},\] (5) where \(\Delta_{\mu}(t)=-\frac{\partial\ell(f_{\mu},y_{\mu})}{\partial f_{\mu}}|_{f_{ \mu}(t)}\) are error signal for each data point.

Abstracts, we can consider the uncoupled variances \(\mathbf{\kappa}\) as "sources" of finite-width noise for each order parameter and the \(\mathbf{D}\) blocks as summarizing a directed causal graph which captures how this noise propagates in the network (through layers and network predictions). In Figure 1, we illustrate this graph showing directed lines that represent causal influences of order parameters on fields and vice versa. For instance, if \(\Phi^{\ell}\) were perturbed, \(D^{\Phi^{\ell+1},\Phi^{\ell}}\) would quantify the resulting perturbation to \(\Phi^{\ell+1}\) through the fields \(h^{\ell+1}\).

In App. D, we calculate \(\mathbf{\kappa}\) and \(\mathbf{D}\) tensors, and show how to use them to calculate the propagator. As an example of our results:

**Proposition 2**: _Partition \(\mathbf{q}\) into primal \(\mathbf{q}_{1}=\text{Vec}\{f_{\mu}(t),\Phi_{\mu\nu}^{\ell}(t,s)...\}\) and conjugate variables \(\mathbf{q}_{2}=\text{Vec}\{\hat{f}_{\mu}(t),\hat{\Phi}_{\mu\nu}^{\ell}(t,s)...\}\). Let \(\mathbf{\kappa}=\frac{\partial^{2}}{\partial\mathbf{q}_{2}\partial\mathbf{q}_{2}^{2}}S[ \mathbf{q}_{1},\mathbf{q}_{2}]\) and \(\mathbf{D}=\frac{\partial^{2}}{\partial\mathbf{q}_{2}\partial\mathbf{q}_{1}^{2}}S[\mathbf{q}_ {1},\mathbf{q}_{2}]\), then the propagator for \(\mathbf{q}_{1}\) has the form \(\mathbf{\Sigma}_{\mathbf{q}_{1}}=\mathbf{D}^{-1}\mathbf{\kappa}\left[\mathbf{D}^{-1}\right]^{\top}\) (App D). The variables \(\mathbf{q}_{1}\) are related to network observables, while conjugates \(\mathbf{q}_{2}\) arise as Lagrange multipliers in the DMFT calculation. From the propagator \(\mathbf{\Sigma}_{\mathbf{q}_{1}}\) we can read off the variance of network observables such as \(N\text{Var}(f_{\mu})\sim\Sigma_{f_{\mu}}\)._

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. E. We proceed to solve the equations defining \(\mathbf{\Sigma}\) 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 \(\gamma\to 0\) where features are frozen. To prevent divergence of the network in this limit, we

Figure 1: The directed causal graph between DMFT order parameters (blue) and fields (green) defines the \(D\) tensors of our theory. Each arrow represents a causal dependence. \(K\) denotes the NTK.
use a background subtracted function \(\tilde{f}(\mathbf{x},\mathbf{\theta})=f(\mathbf{x},\mathbf{\theta})-f(\mathbf{x},\mathbf{\theta}_{0})\) which is identically zero at initialization [4]. For mean square error, the \(N\rightarrow\infty\) and \(\gamma\to 0\) limit is governed by \(\frac{\partial\tilde{f}(\mathbf{x})}{\partial t}=\mathbb{E}_{\mathbf{x}^{\prime} \sim\mathcal{D}}\Delta(\mathbf{x}^{\prime})K(\mathbf{x},\mathbf{x}^{\prime})\) with \(\Delta(\mathbf{x})=y(\mathbf{x})-\tilde{f}(\mathbf{x})\) (for MSE) and \(K\) is the static NTK. The finite-\(N\) initial covariance of the NTK has been analyzed in prior works [27, 13, 14], which reveal a dependence on depth and nonlinearity. Since the NTK is static in the \(\gamma\to 0\) limit, it has constant initialization variance through training. Further, all sensitivity blocks of the Hessian involving the kernels and the prediction errors \(\mathbf{\Delta}\) (such as the \(D^{\Phi^{\prime},\Delta}\)) vanish. We represent the covariance of the NTK as \(\kappa(\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{2},\mathbf{x}_{3})=N\text{Cov}(K(\mathbf{x}_{1}, \mathbf{x}_{2}),K(\mathbf{x}_{3},\mathbf{x}_{4}))\). To identify the dynamics of the error \(\mathbf{\Delta}\) covariance, we consider the eigendecomposition of the infinite-width NTK \(K_{\infty}(\mathbf{x},\mathbf{x}^{\prime})=\sum_{k}\lambda_{k}\psi_{k}(\mathbf{x})\psi_{k }(\mathbf{x}^{\prime})\) with respect to the training distribution \(\mathcal{D}\), and decompose \(\kappa\) in this basis

\[\kappa_{k\ell mn}=\left\langle\kappa(\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{3},\mathbf{x}_ {4})\psi_{k}(\mathbf{x}_{1})\psi_{\ell}(\mathbf{x}_{2})\psi_{n}(\mathbf{x}_{3})\psi_{m}( \mathbf{x}_{4})\right\rangle_{\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{3},\mathbf{x}_{4}\sim \mathcal{D}}, \tag{6}\]

where averages are computed over the training distribution \(\mathcal{D}\).

**Proposition 3**: _For MSE loss, the prediction error covariance \(\mathbf{\Sigma}^{\Delta}(t,s)=N\text{Cov}_{0}(\mathbf{\Delta}(t),\mathbf{\Delta}(s))\) satisfies a differential equation (App. G)_

\[\left(\frac{\partial}{\partial t}+\lambda_{k}\right)\left(\frac{\partial}{ \partial s}+\lambda_{\ell}\right)\Sigma_{k\ell}^{\Delta}(t,s)=\sum_{nm}\kappa _{km\ell n}\Delta_{n}^{\infty}(t)\Delta_{n}^{\infty}(s), \tag{7}\]

_where \(\Delta_{k}^{\infty}(t)\equiv\exp\left(-\lambda_{k}t\right)\left\langle\psi_{k }(\mathbf{x})y(\mathbf{x})\right\rangle_{\mathbf{x}}\) are the errors at infinite width for eigenmode \(k\)._

An example verifying these dynamics is provided in App. Fig. A.1. 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 corrections rather than from initialization variance (App. G.1). Corrections to the mean \(\left\langle\Delta\right\rangle\) are analyzed in App. G.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.

### 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 \(\mathbf{\Delta}\). To gain insight, we first study a simplified setting where the data distribution is a single training example \(\mathbf{x}\) and single test point \(\mathbf{x_{\star}}\) in a two layer network. We will track \(\Delta(t)=y-f(\mathbf{x},t)\) and the test prediction \(f_{\star}(t)=f(\mathbf{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. D).

**Proposition 4**: _Computing the Hessian of the DMFT action and inverting (App. H), 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 }_{+}}\)_

\[\mathbf{\Sigma}_{\mathbf{q}_{1}}=\begin{bmatrix}\mathbf{I}+\mathbf{\Theta}_{K}&0&\mathbf{ \Theta}_{\Delta}&0\\ -\mathbf{\Theta}_{K_{\star}}&\mathbf{I}&0&-\mathbf{\Theta}_{\Delta}\\ -\mathbf{D}&0&\mathbf{I}&0\\ -\mathbf{D}_{\star}&0&0&\mathbf{I}\end{bmatrix}^{-1}\begin{bmatrix}0&0&0&0\\ 0&0&0&0\\ 0&0&\mathbf{\kappa}&\mathbf{\kappa}_{\star\star}^{\top}\\ 0&0&\mathbf{\kappa}_{\star}&\mathbf{\kappa}_{\star\star}\end{bmatrix}\begin{bmatrix} \mathbf{I}+\mathbf{\Theta}_{K}&0&\mathbf{\Theta}_{\Delta}&0\\ -\mathbf{\Theta}_{K_{\star}}&\mathbf{I}&0&-\mathbf{\Theta}_{\Delta}\\ -\mathbf{D}&0&\mathbf{I}&0\\ -\mathbf{D}_{\star}&0&0&\mathbf{I}\end{bmatrix}^{-1\top}, \tag{8}\]

_where \([\mathbf{\Theta}_{K}](t,s)=\Theta(t-s)K(s)\), \([\mathbf{\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\) quantify sensitivity of the kernel to perturbations in the error signal \(\Delta(s)\). Lastly \(\kappa,\kappa_{\star},\kappa_{\star\star}\) correspond to uncoupled variances (and covariance) for \(\{K(t),K_{\star}(t)\}\)._
Figure 2: 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{\mathbf{x}\cdot\mathbf{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.
In Fig. 2, we plot the resulting theory (diagonal blocks of \(\mathbf{\Sigma}_{\mathbf{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 \(\left\langle K\right\rangle\) and test predictions \(\left\langle f_{\star}\right\rangle\) change by a larger amount for larger \(\gamma\) (equation (64)). 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\to 0\) the curve approaches \(N\) 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.2, 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.

### Offline Training with Multiple Samples

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. D). 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}(\frac{P}{N})\). In the case of two-layer linear networks trained on whitened data (\(\frac{1}{D}\mathbf{x}_{\mu}\cdot\mathbf{x}_{\nu}=\delta_{\mu\nu}\)), the equations for the propagator simplify and one can separately solve for the variance of \(\mathbf{\Delta}(t)\in\mathbb{R}^{P}\) along signal direction \(\mathbf{y}\in\mathbb{R}^{P}\) and along each of the \(P-1\) orthogonal directions (App. I). At infinite width, the task-orthogonal component \(\mathbf{\Delta}_{\perp}\) vanishes and only the signal dimension \(\Delta_{y}(t)\) evolves in time with differential equation [9, 46]

\[\frac{d}{dt}\Delta_{y}(t)=2\sqrt{1+\gamma^{2}(y-\Delta_{y}(t))^{2}}\ \Delta_{y}(t)\,\ \mathbf{\Delta}_{\perp}(t)=0. \tag{9}\]

However, at finite width, both the \(\Delta_{y}(t)\) and the \(P-1\) orthogonal variables \(\mathbf{\Delta}_{\perp}\) inherit initialization variance, which we represent as \(\Sigma_{\Delta_{y}}(t,s)\) and \(\Sigma_{\perp}(t,s)\). In Fig. 3 (a)-(b) we show this approximate solution \(\left\langle|\mathbf{\Delta}(t)|^{2}\right\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 I 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. 3 (b) shows that as \(P\) increases, the leading order approximation breaks down as higher order terms become relevant.

## 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. D). 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^{\ell-1}_{\alpha\beta}(t^{\prime},s^{\prime})}\left\langle\phi( \hbar_{\mu}^{\ell}(t))\phi(\hbar_{\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. K. An example set of results for a depth 4 network is provided in Fig. 4. 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\text{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.5.

## 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
Figure 4: 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 \(\left\langle K(t)\right\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\).

Figure 3: Large input dimension or multiple samples amplify finite size effects in a simple two layer model. 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.
[19, 20, 21]. In these settings, the kernel undergoes an initial scale 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 \(|\mathbf{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. L)

\[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. 5 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, F.2) In Fig. 6, 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. 6(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. F.1). Our analysis in App. F.1 suggests that the rate of convergence receives a finite size correction with leading correction \(\mathcal{O}(N^{-1})\) F.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. 6(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\) on CIFAR-10.

In App. Fig. A.6, 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.7, 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.

Figure 5: Edge of stability effects do not imply deviations from infinite width behavior. Black dashed lines are theory. (a) The loss dynamics for width \(N=500\) networks (solid colors) compared to the infinite width DMFT (dashed black). (b) The average kernel over an ensemble of several NNs. 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\). (c)-(d) Remarkably variance due to finite size can reduce during training on both sides of the edge of stability (for \(\gamma\) smaller and larger than the critical value \(\sim 1/\eta\)), suggesting that infinite width theory can be predictive of wide but finite networks even in the large learning rate regime.
Figure 6: 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. F) 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.
Discussion

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 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, the leading term involving only \(\mathbf{\Sigma}\) does not capture the complete finite size distribution defined in Eq. (3), especially as the sample size becomes comparable to the width. Future work could explore in greater detail the higher order contributions from averages involving powers of \(V(\mathbf{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]. Further, even though we expect our perturbative expressions to give a precise asymptotic description of finite networks in mean field/\(\mu\)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 or high dimensional limits where the analysis may simplify.

### Acknowledgements

This work was supported by NSF Award DMS-2134157. BB thanks Alex Atanasov, Jacob Zavatone-Veth, Boris Hanin and Greg Yang for helpful discussions and for their comments on this manuscript. BB also acknowledges Jeremy Cohen for discussions about large step size dynamical effects and their relationship to network width.
