SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In [12, 15], gives an equivalent definition of  $\mu \mathrm{P}$  that gives rise to the same dynamics but keeps the learning rate to be  $\Theta (1)$ . We use this version of  $\mu \mathrm{P}$  in the experiments that we run, simply because that is what is used in the package [13]. Consequently, our learning rate does not need to be changed as width grows.
D Importance of Parameterization
Figure 13: (a-b) An experiment on CIFAR-10 with a depth 4 CNN (no layernorm) showing that the finite width consistency at the range of widths  $N$  considered in this work is special for  $\mu \mathrm{P}$  and is not observed in NTK parameterization. Width-consistent kernel regime behavior requires a much larger width. Networks are parameterized to agree at  $N = 128$ . We see that in NTK parameterization wider networks train more slowly and their kernels align less to the target function. The reason for the stronger consistency of  $\mu \mathrm{P}$  networks is that feature learning remains  $\Theta(1)$ .
In this section, we aim to illustrate that many of the claims of width consistency at practical widths are in fact contingent on using the mean field/  $\mu$ -parameterization. In Figure 13 we compare the early dynamics of network training on CIFAR-10 between networks in the NTK parameterization (dashed lines) and in mean field parameterization (solid). The models are initialized in such a way so that their  $N = 128$  dynamics perfectly coincide. We plot both test error and the kernel alignment between the NTK and the target function. While both parameterizations will eventually converge to an infinite width limit, we note the following two differences in width scaling behavior:
1. Wider networks tend to train more slowly in NTK parameterization and perform worse as the width grows. This is because the rate of feature learning is not held constant across widths and decreases as width grows. Wider networks in  $\mu \mathrm{P}$  train slightly faster.
2. Mean field networks approach their limit more quickly as width increases in terms of both test loss and kernel dynamics.
The general theoretical principle behind these observations is that mean field parameterization generates feature updates which are scale independent, thus eliminating an unnecessary source of finite size approximation error in the dynamics [39].
E Further Studies of  $\mu \mathbf{P}$  versus SP
E.1 MLPs
In figure 14, we show a 3-layer MLP learning a quadratic polynomial. In subfigure a) use a batch size of 10 and a learning rate going as  $\eta = 5N / (1 + \alpha_0^2)$  with  $\alpha_0 = 1$ . The output layer is scaled as  $\alpha_0 / \sqrt{N}$ , putting us in the rich regime. The learning rate has been picked to be nearly as large as possible at this batch size in order to maximize the large loss curve fluctuations yielded by large learning rate effects. In subfigure c) we do not rescale the output layer, and have a width-independent learning rate going as  $\eta = 50 / (1 + \alpha_0^2)$  with  $\alpha_0 = 1$ . This puts is in the large-learning rate regime for a standard parameterized network. See Appendix A for more details.
We plot the learning curves across widths and find striking agreement, even at the fine-grained level of fluctuations due to small batch size and large learning rate effects. Although this is not exactly the full-batch edge-of-stability effect reported in [7], the large oscillations may be similar to a small batch size analog. We plot the absolute difference from the widest network in subfigure b) to highlight the strong agreement across widths.
In subfigure c), we have the same network but in standard parameterization. The narrower networks now learn features more quickly, leading to inconsistent dynamics across widths.
Figure 14: 3-Layer MLP learning a quadratic polynomial  $y = Q_{2}(\boldsymbol{\beta} \cdot \boldsymbol{x})$  on the sphere. Data is provided in an online setting in the same order across widths, as in the realistic experiments in the main text. Large learning rate small batch size effects in MLPs are consistent across large widths. a) For  $\mu \mathrm{P}$  networks, the loss curves match across widths, even accounting for fluctuations due to batch size or large learning rate edge-of-stability-like effects. b) Plot of the difference in training error from the widest network. c) The same network but in standard parameterization. Dynamics are no longer consistent across widths, and wider networks approach a lazy limit.
E.2 Vision
Next, we focus on a vision task and compare the large learning rate small batch size effects in SP to the  $\mu \mathrm{P}$  parameterized network in Figure 6a. By contrast to that figure, we see significantly different dynamics and batch variation across widths. In Figure 15 a) we plot the early time behavior of a CIFAR-5m task at large learning rate. The large learning rate effects cause the loss to substantially oscillate, but the oscillations across widths are inconsistent by contrast to 6a b). Further, at late times in Figure 15, the sharp spikes in the loss function due to large learning rate effects become substantially different across widths. Indeed, in SP some widths converge for a given learning rate while others do not. This trend has already been well-studied in [15]. We again stress that our observation is that not only are the final losses similar across widths in  $\mu \mathrm{P}$  (as observed in [15]), but that the individual batch and large learning rate fluctuations agree across widths at early times in  $\mu \mathrm{P}$  as well.
Figure 15: Different widths have different loss curves. a) Early time dynamics of the loss across widths is not consistent. b) Dynamics of the loss across widths at later times also does not appear consistent. There are explosions that happen at different times and scales across widths.
E.3 Language
Finally, we present a complementary set of figures to those in the right columns of Figures 3 and 12 for transformers of the same architecture on Wikitext-103 but in standard parameterization.
(b) Correct logit value
Figure 16: An analog of Figure 3: SP transformers trained with Adam. a) Training loss. At large widths, the learning rate chosen is too big for the network to properly learn, and the loss is flat. This is consistent with what is observed in [15] — the optimal learning rate in SP changes with width. b) Value placed on the correct logit for a specific masked token. c) RMSE of correct logit value from the widest network. In both of these plots, the monotonic behavior across width evident under the  $\mu$ P parameterization is violated. Even after discarding the networks that do not converge under SP, the behavior remains non-monotonic across width.
(c) RMSE from widest
F Bias-Variance Decompositions over Initializations
In this section we explicitly define what we mean by initialization bias and initialization variance. Following [30], we consider the trained neural network function  $f_{\theta^{*}}(\pmb{x})$  to depend on the training set  $\mathcal{D}$  (including inputs  $\pmb{x}$ , outputs  $y$ , and possible label noise  $\epsilon$ ) as well as the initial parameters  $\theta_0$ . Here,  $\theta_0$  are the initial parameters and  $\theta^{*}$  are the final parameters, which are implicitly functions of the initial ones.
Classical statistical learning theory often focuses on the variance of the learned function as a function of the training samples given  $(\boldsymbol{x}_{\mu},y_{\mu})_{\mu = 1}^{P}$ . By contrast, our paper focuses on the variance due to the initial parameters  $\theta_0$  from which training begins. In the overparameterized regime with more parameters than data-points and relatively little label noise, this has been shown to be the dominant source of variance for neural networks [27, 29, 28].
Consequently, our definition of the bias of a neural network is given by
$$
\bar {f} (\boldsymbol {x}) = \mathbb {E} _ {\theta_ {0}} \left[ f _ {\theta^ {*} \left(\theta_ {0}\right)} (\boldsymbol {x}) \right]. \tag {8}
$$
The bias can be approximated by averaging a sufficiently large ensemble of neural networks over initialization seeds. Each network in the ensemble is trained on the same dataset in the same batch order using the same optimizer.
The variance of the neural network predictor is given by:
$$
\operatorname {V a r} _ {\theta_ {0}} f = \mathbb {E} _ {\theta_ {0}} \left[ \left(f _ {\theta^ {*} (\theta_ {0})} (\boldsymbol {x}) - \bar {f} (\boldsymbol {x})\right) ^ {2} \right]. \tag {9}
$$
The variance of  $E$  ensembles of a given network with independent initialization seeds  $\theta_{e}$  is given by
$$
\operatorname {V a r} _ {\theta_ {0}} \left[ \frac {1}{E} \sum_ {e} f _ {\theta^ {*} (\theta_ {e})} (\boldsymbol {x}) \right] = \mathbb {E} _ {\theta_ {0}} \left[ \left(\frac {1}{E} \sum_ {e} f _ {\theta^ {*} (\theta_ {e})} (\boldsymbol {x}) - \bar {f} (\boldsymbol {x})\right) ^ {2} \right]. \tag {10}
$$
As  $E \to \infty$ , the empirical average of  $E$  network ensembles approaches the bias, so the variance of the ensembled networks goes to zero. As long as the errors are uncorrelated, the variance the network ensembles decreases as  $1 / E$ . In practice, we see that comparable decay rates are achieved.
G Task-Dependent Scaling Laws in Width
For the settings where we observed larger deviations in the dynamics for models of varying widths, we examined scaling of the training losses with respect to width  $N$  after a significant amount of training (2 epochs for ImageNet and 1 epoch for Wikitext 103). We fit power laws of the form  $\mathrm{Loss} = CN^{-\beta} + D$  where  $C,\beta ,D$  are fit to the data using the 'scipy适合自己' function. The resulting fits are provided in Figure 17. We find an excellent power law fit, with  $R^2$  above .99.
The existence of such power law behavior across widths provides further evidence of the networks approaching a well-defined infinite width limit. One could imagine that the differences between successive networks might get smaller (as in Figure 4) but that there is no well-defined limit as  $N \to \infty$ , similar to the terms in a harmonic sum. The fact that the power law fit has exponent much larger than 0 and does not display logarithmic dependence on  $N$  provides empirical evidence that we expect convergene as  $N \to \infty$ .
Additionally, given that the observed power laws  $N^{-\beta}$  are task-dependent and significantly different from  $N^{-1}$  suggests that models at late time are not well described by perturbation theory around the infinite width mean field limit, which predicts a universal exponent of  $\beta = 1$  [54, 39]. This motivates novel theoretical descriptions of finite width mean field
Figure 17: Train losses of Wikitext and Imagenet models are well-fit by power laws.
learning networks at late time which can capture task-dependent exponents.
H Overview of Finite Width Corrections to Feature Learning Networks
In this section we review some basic ideas from the mean field theory of feature learning neural networks. We first describe the predictions that mean field theory makes about infinite width networks before describing finite size corrections to the dynamics of learning. To eliminate unnecessary complexity, we will focus on MLP layers, but these arguments can be easily extended to CNN and self-attention layers as well. We start by defining a MLP in a parameterization equivalent to  $\mu \mathrm{P}$
$$
f _ {\mu} = \frac {1}{N} \boldsymbol {w} ^ {L} \cdot \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right), \boldsymbol {h} _ {\mu} ^ {\ell + 1} = \frac {1}{\sqrt {N}} \boldsymbol {W} ^ {\ell} \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right), \boldsymbol {h} _ {\mu} ^ {1} = \frac {1}{\sqrt {D}} \boldsymbol {W} ^ {0} \boldsymbol {x} _ {\mu}. \tag {11}
$$
We will consider these networks trained from a random Gaussian initialization of the weights so that  $\pmb{\theta} = \mathrm{Vec}\{\pmb{w}^L, \dots, \pmb{W}^0\}$  follows  $\pmb{\theta} \sim \mathcal{N}(0, I)$  at initialization. This network is then trained with some gradient based optimizer, leading to dynamical predictions  $f_{\mu}(t)$  and dynamical preactivations  $\pmb{h}_{\mu}^{\ell}(t)$ . Because of the random initialization of weights, the outputs of the network and the precise preactivations are random variables. However, at infinite width  $N \to \infty$ , a dramatic simplification of the dynamics occurs.
H.1 The Infinite Width/Mean Field Limit
The predictions  $f_{\mu}(t)$  and internal representations of infinite width limit of neural networks admit a description in terms of non-random initialization-independent dynamical feature kernels  $\Phi_{\mu \nu}^{\ell}(t,s)$  and gradient kernels  $G_{\mu \nu}^{\ell}(t,s)$  defined as
$$
\Phi_ {\mu \nu} ^ {\ell} (t, s) = \frac {1}{N} \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell} (t)\right) \cdot \phi \left(\boldsymbol {h} _ {\nu} ^ {\ell} (s)\right), G _ {\mu \nu} ^ {\ell} (t, s) = \frac {1}{N} \boldsymbol {g} _ {\mu} ^ {\ell} (t) \cdot \boldsymbol {g} _ {\nu} ^ {\ell} (s), \tag {12}
$$
where  $\mu, \nu$  index data points and  $t, s$  index training time and  $\pmb{g}_{\mu}^{\ell}(t) = N\frac{\partial f_{\mu}}{\partial\pmb{h}^{\ell}}$  are back-propagated gradient signals [12, 14]. Further, all preactivation vectors  $\pmb{h}_{\mu}^{\ell}(t) \in \mathbb{R}^{N}$  have entries that become iid draws from a (potentially non-Gaussian) single site density  $p(h)$ , which converges as
$$
\frac {1}{N} \sum_ {i = 1} ^ {N} \delta \left(h - h _ {i}\right)\rightarrow p (h), \tag {13}
$$
which should be understood in terms of integration of these densities against test functions. At infinite width, the sums over neurons in a layer can be replaced by deterministic integrals over this single site density  $\Phi_{\mu \nu}^{\ell}(t,s) = \int p(h_{\mu}^{\ell}(t),h_{\nu}^{\ell}(s))\phi (h_{\mu}^{\ell}(t))\phi (h_{\nu}^{\ell}(s))dh_{\mu}^{\ell}(t)dh_{\nu}^{\ell}(s)$ .
H.2 Finite Width Effects
At finite width, the internal kernels  $\{\Phi_{\mu \nu}^{\ell}(t,s),G_{\mu \nu}^{\ell}(t,s)\}$  and predictions  $f_{\mu}(t)$  of the model become initialization and width-dependent and deviate from their mean field dynamics. For Gaussian random initialization of the weights of the network, the predictions and kernels fluctuate (from init to init) with variance that scales asymptotically like  $\Theta (1 / N)$  for width  $N$  (or  $1 / d_{model}$  for transformer) [39]. Further, the ensemble averaged values for the predictions  $\langle f_{\mu}(t)\rangle$  and kernels  $\langle \Phi_{\mu \nu}^{\ell}(t,s)\rangle$  differ asymptotically from their infinite width values by  $\Theta (N^{-1})$ . Both of these two leading order effects can influence the expected (train or test) loss of the model. At fixed width and late training time, finite size effects beyond leading order can accumulate and become relevant, however theory predicts that any observable average at width  $N$  admits an asymptotic series in powers of  $N^{-1}$  [39].
H.2.1 Trainability at Finite Size
The  $\Theta (N^{-1})$  correction to feature and gradient kernels can lead to non-trivial corrections to the loss dynamics. Working in continuous time, we can define the neural tangent kernel (NTK) as  $K_{\mu \nu}(t) = \sum_{\ell}G_{\mu \nu}^{\ell +1}(t,t)\Phi_{\mu \nu}^{\ell}(t,t)$ , where base cases are  $\Phi_{\mu \nu}^{0}(t,s) = \frac{1}{D}\pmb{x}_{\mu}\cdot \pmb{x}_{\nu}$  and  $G_{\mu \nu}^{L + 1}(t,s) = 1$ . Following the approximation to online dynamics with MSE loss in Section 4, we consider a gradient flow on the average dynamical NTK
$$
\frac {d}{d t} \boldsymbol {\Delta} (t) = - \left\langle \boldsymbol {K} (t) \right\rangle \boldsymbol {\Delta} (t) \Rightarrow \boldsymbol {\Delta} (t) = \mathcal {T} \exp \left(- \int_ {0} ^ {t} d s \left\langle \boldsymbol {K} (s) \right\rangle\right) \boldsymbol {y}, \tag {14}
$$
where  $\mathcal{T}$  is the time-ordering operator. We now consider the leading correction to the average NTK around infinite width  $\langle K(t)\rangle = K_{\infty}(t) + \frac{1}{N} K^{1}(t) + \Theta (N^{-2})$ . With this correction, we see that the dynamics of errors  $\pmb{\Delta}$
$$
\boldsymbol {\Delta} (t) = \mathcal {T} \exp \left(- \int_ {0} ^ {t} d s \boldsymbol {K} _ {\infty} (s) - \frac {1}{N} \int_ {0} ^ {t} d s \boldsymbol {K} ^ {1} (s) + \Theta (N ^ {- 2})\right) \boldsymbol {y}. \tag {15}
$$
The fact that the  $\frac{1}{N} K^1$  correction is integrated over time and placed in the matrix exponential indicates that small corrections to NTK dynamics can lead to large dynamical amplification of logit corrections. This fact was pointed out in another work [39] which tried to motivate a study of perturbation theory in logarithms of the transition matrix  $\log T(t)$  defined as
$$
\frac {d}{d t} \boldsymbol {T} (t) = - \langle \boldsymbol {K} (t) \rangle \boldsymbol {T} (t), \boldsymbol {T} (0) = \boldsymbol {I}, \boldsymbol {R} (t) = \log \boldsymbol {T} (t). \tag {16}
$$
The solution to this can be used to construct the errors at a later time  $\pmb{\Delta}(t) = \exp(\pmb{R}(t))\pmb{y}$ .
I The empirical sufficiency of ensembling only a few times
Our experiments rely on ensembles of small numbers of neural networks to analyze the bias component of the loss as it varies across width. One can show the marginal value of adding a network to the ensemble decreases with the ensemble size. Figure 18 illustrates this in the setting of ResNets trained on ImageNet. In Figure 18(a), we show that the reduction in variance over initializations due to ensembling rapidly plateaus as soon as the ensemble size reaches  $E = 3$ . In Figure 18(b), we show that the loss curves as function of width are very similar for ensemble sizes above 3. Indeed, the green, orange, and red curves — corresponding to  $E \in \{3,4,5\}$  — are nearly identical.
Lastly, Figure 18(c) confirms that the initialization variance plays a negligible role in the scaled RMSE distances between true logits across width, even for single draws of networks trained on a small number of examples. The dashed lines in this Figure correspond to scaled RMSE distances between networks of the same widths as their solid-line companions, albeit with weights initialized to zero. The ordering and scale of the curves are comparable, and in fact nearly identical for the curves corresponding to  $N \in \{128, 256\}$  compared to  $N = 512$ .
(a) Test loss vs. # ensembles
(b) Test loss vs. width
Figure 18: ResNet-18/ImageNet loss as a function of a) ensemble size or b) width after 2 epochs of training with heavy data augmentation. c) Comparing original network with a network whose output has been set to zero at the start of training.
(c) Test loss for subtracted networks
J Offline Training
Figure 19 depicts the loss curve for a ConvNeXt-T (tiny) model trained on ImageNet in the typical, offline setting — where data is encountered repeatedly across many epochs. As the networks overfit the training data — in Figure 19, beyond 40,000 training steps or five epochs — the loss curves diverge dramatically for different-width networks. Width consistency subsequently erodes.
Figure 19: Consistency of loss curves across widths in the beginning and separation as loss becomes sufficiently small in offline learning.
K Use of Compute
For most experiments, we used Nvidia A100 SXM4 40GB and 80 GB GPUs on an academic cluster.
For the Wikitext-103 tasks, each width included 4 ensembles loaded onto an A100 GPU that ran for a range between 1 to 3 days. For each sweep over widths this corresponds to about 8 A100-days. Accounting for sweeps over different sequence lengths, optimizers, and parameterizations, this corresponds to about 50 A100-days.
All MLP tasks, including the calculation of empirical NTKs and their spectral properties were done in 15-30 minute Colab sessions using the basic GPUs provided.
The CIFAR-10 ResNet experiments in Figure 10 were done using a total of less than 1 A100-day of compute across all widths and ensembles.
The ImageNet ResNet experiments vectorize training over between one to four same-width neural networks on one A100 GPU. Each experiment training a collection of networks for 30 epochs takes between one to three A100-days. Overall, these experiments expended roughly 30 A100-days.
For the CIFAR-5m experiments in Figure 2 and 3, across all widths, it required a few hours of A100 GPU. For Figure 7 and 8, as these were ensembled across multiple runs, these required close to 1-2 days of A100-GPUs. Figure 6a was just run for a few 100 steps of the training, so didn't use much compute power.