SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(f) Noisy CIFAR-50k test error
Figure 8: Top Row: Effects of offline training on train metrics. In (a) and (b) we do multi-epoch training on CIFAR-200k. We see that both bias and variance for train error are magnified by offline training and do not tend to 0 for the largest widths we could try. (c) we do multi-epoch training on noisy CIFAR-50k and again observe large bias and variance terms at large widths. Bottom Row: Effects of offline training on test metrics. In (d) and (e) we do multi-epoch training on CIFAR-200k. We see that both bias and variance for test error are near 0 at large widths. (f) We train on noisy CIFAR-50k and observe that "wider is better" is violated for ensembled networks.
deviations in width for single network performance, indicating that the benign overfitting observed in Figure 8 (d) and (e) is dataset-dependent. Furthermore, variance is found to be much larger than in the non-noisy experiments.
- Surprisingly, we discover (Figure 8 (f)) that some ensembled narrower width networks outperform ensembled wider networks. This presents a counterexample to the "wider is better" phenomenon [15] for ensembled networks. We hypothesize that such counterexamples can only exist in the context of offline training.
4 Spectral perspective on finite-width bias
In this last section, we develop a toy model in which the effect of finite-width bias can be clearly seen. We analyze it first in the simple setting of an MLP fitting a polynomial in the lazy limit. Here, all the dynamics are well-captured by the finite-width empirical neural tangent kernel (eNTK). By studying the spectral properties of this kernel across widths, we see that finite widths lead to eNTK's with worse bias components in their losses.
Concretely, we see that although the eigenvalue spectrum of the ensembled eNTK is not substantially affected by finite width, the decomposition of the task into eNTK eigenvectors changes, with narrow-width eNTK's putting more of the task into smaller eigenmodes that take longer to be learned. In practice, applying this analysis to the after-kernel of the trained ResNets on CIFAR-5m reveals similar behavior. Prior literature has demonstrated that many of the properties of the final learned function are captured by the after-kernel [50-52].
We consider a model of online learning where a large batch of data from the population distribution  $p(\pmb{x})$  is sampled at each step. This leads to approximate gradient flow dynamics  $\frac{d}{dt}\pmb{\theta} = -\frac{1}{2}\nabla_{\theta}\mathbb{E}_{\pmb{x}}(f(\pmb{x},\pmb{\theta}) - y(\pmb{x}))^2$  (Appendix H). To analyze this equation, we choose a fixed orthonormal basis  $\{\psi_k(\pmb{x})\}$  for the space  $L^2 (\mathbb{R}^D,p(\pmb{x})dx)$  of square-integrable functions on input space. The function  $f(x)$ , residual error  $\Delta (\pmb {x}) = y(\pmb {x}) - f(\pmb {x})$ , and the empirical NTK
$K(\pmb{x},\pmb{x}^{\prime},t)$  can be expressed in this basis as  $f(\pmb {x},t) = \sum_{k}f_{k}\psi_{k}(\pmb {x}),\Delta (\pmb {x},t) = \sum_{k}\Delta_{k}\psi_{k}(\pmb {x})$  and  $K(\pmb {x},\pmb{x}^{\prime},t) = \sum_{k\ell}K_{k\ell}(t)\psi_{k}(\pmb {x})\psi_{\ell}(\pmb{x}^{\prime})$  , respectively. Their training evolution is given by:
$$
\frac {d}{d t} f (\boldsymbol {x}, t) = \mathbb {E} _ {\boldsymbol {x} ^ {\prime} \sim p (\boldsymbol {x})} K \left(\boldsymbol {x}, \boldsymbol {x} ^ {\prime}, t\right) \Delta \left(\boldsymbol {x} ^ {\prime}, t\right) = - \sum_ {k \ell} K _ {k \ell} (t) \Delta_ {\ell} (t) \psi_ {k} (\boldsymbol {x}). \tag {1}
$$
The statistics of the dynamical NTK matrix  $K_{kl}(t)$  summarizes the statistics of the error dynamics  $\Delta(\boldsymbol{x}, t)$  at any level of feature learning. At infinite width,  $K_{k\ell}(t)$  is deterministic, while at finite width, it receives a  $\Theta(N^{-1})$  mean displacement and a  $\Theta(N^{-1/2})$  fluctuation around its mean [24, 25, 39]. We consider approximating the dynamics of the ensemble predictor by  $\frac{d}{dt} \langle f_k(t) \rangle_{\theta_0} \approx \sum_{\ell} \langle K_{k\ell}(t) \rangle \langle \Delta_\ell(t) \rangle$ . Here,  $\langle \cdot \rangle$  denotes averages over initializations. This expression neglects the contribution from  $\mathrm{Cov}(K_{k\ell}, \Delta_\ell)$ . We show that this approximation is accurate in depth-3 MLPs trained on Gegenbauer polynomial regression tasks in Figure 9 a). For more details see Appendix A.
In the lazy limit, the kernel is static and we choose  $\psi_{k}$  to diagonalize  $\langle K_{k\ell}\rangle = \delta_{k\ell}\lambda_k$ . This yields the loss dynamics  $\mathcal{L}(t) = \sum_{k}\langle y(\pmb {x})\psi_{k}(\pmb {x})\rangle^{2}e^{-2\lambda_{k}t}$ . We can therefore quantify alignment of eigenfunctions to task with the cumulative power distribution  $C(k) = \sum_{\ell < k}\langle y(\pmb {x})\psi_{\ell}(\pmb {x})\rangle_{\pmb{x}}^{2} / \langle y(\pmb {x})^{2}\rangle_{\pmb{x}}$  which is the proportion of the task that is explained by the first  $k$  eigenvectors [53]. If  $C(k)$  rises rapidly with  $k$  then the loss falls faster [53]. In this limit, there are two ingredients that could make the bias dynamics across widths distinct. First, the eigenvalues  $\lambda_{k}$  which set the timescales could be width-dependent. Second, the eigenfunctions  $\psi_{k}(\pmb {x})$  that diagonalize  $\langle K\rangle$  can change with width. In Figures 9 b) and c) we show that the dominant effect is the latter. Finite-width corrections do not substantially affect the spectrum, but they do increase the proportion of the target function that lies in the eigenspaces corresponding to modes that are slower to learn.
To test whether these findings continue to hold in more realistic experiments, we computed the final NTKs (after kernels) of the ResNet-18 models trained on CIFAR-5M (specifically the models from Figures 3, 5). We ensemble average to get kernel  $\langle K_{c,c'}(\pmb{x},\pmb{x}')\rangle$  for output channels  $c,c'$  and input images  $\pmb{x},\pmb{x}'$ . We then compute the kernel gradient flow corresponding to MSE training on the true target function for CIFAR-5M  $\frac{d}{dt}\Delta_c(\pmb{x}) = -\sum_{c'}\mathbb{E}_{\pmb{x}'}\langle K_{c,c'}(\pmb{x},\pmb{x}')\rangle \Delta_{c'}(\pmb{x}')$  from initial condition given by the one-hot target labels  $\Delta_c(\pmb{x})|_{t=0} = y_c(\pmb{x})$ . The convergence rate of this dynamical system is again set by the eigenvalues and eigenfunction-task alignment. In Figure 9 (d), we find that the after kernels for wider networks give slightly more rapid convergence. Figures 9 (e) and (f) show that, similar to the MLP experiment, the spectra are very consistent across widths, but the eigenfunction task alignments, measured with  $C(k)$  are not. Overall, these experiments suggest that an important aspect of the bias of finite width models compared to their infinite width analogs is the deformation of their eigenfunctions.
We have demonstrated a striking consistency across widths for many quantities of interest to deep learning practitioners. Our fine-grained studies go beyond simply comparing test losses and have demonstrated that learned network functions, internal representations, and dynamical large learning rate phenomena agree for sufficiently large widths on a variety of tasks across vision and language. At later training times, or after many repetitions of the dataset, we observe systematic deviations brought on by finite width, and have characterized them in terms of the bias and variance of the network over initializations. This study motivates the applicability of infinite-width feature-learning models (and the accumulating finite width deviations from this limit) in reasoning about large scale models trained on real-world data.
In light of the accumulation of finite-width deviations at later training times, we caution that our study only exhibits the consistency of the infinite-width limit with the training time held fixed. It does not make a claim about other possibly limits that vary the training time jointly with the width to infinity, perhaps along a compute frontier. We leave further inquiry into such limits for future work.
(a) MLP Online training losses
(b) MLP Kernel Spectra
(c) MLP unexplained task fraction
(d) CIFAR-5M After-Kernel
(e) After-Kernel Spectra
(f) After-Kernel unexplained task fraction
Figure 9: Spectral properties of the NTK can account for bias gaps across widths. (a) Depth 3 MLPs in the lazy limit  $(\gamma_0^{-1} = 200)$  learning a quadratic polynomial from a uniform distribution on the sphere in  $D = 5$  dimensions online. Wider networks perform better (dots). Even after ensembling (dashed), wider is better, and the ensembled curves match those of the averaged eNTK (solid). (b) The spectra of the averaged eNTK across widths do not show substantial variability. (c) However, at narrower width, the eigen-decomposition of the task has greater weight along higher spectral modes, consequently leading to a slow-down in training. These results hold across dimensions, batch sizes, task complexity, and architectures. Strong feature learning can reduce this effect. See Appendix H (d) We computed the ensemble averaged after kernels from the CIFAR-5M ResNet-18 models and computed the theoretical kernel flow on the task. Wider models have a slightly better mean kernel for this task. (e) The eigenvalues of the final NTKs are very consistent across widths. (f) The eigenfunction-target alignment of the final kernels noticeably differ across widths, evidenced by the cumulative power distribution  $C(k)$  which accounts for the gap in theoretical loss curves under kernel flow.
We thank Boaz Barak, Jeremy Cohen, Alex Damian, Nikhil Ghosh, Gal Kaplun, Eric Michaud, Jamie Simon and Jacob Zavatone-Veth for helpful discussions throughout this project. We also thank Jacob Zavatone-Veth for comments on the draft.
AA is supported by the Professor Yaser S. Abu-Mostafa Fellowship from the Fannie and John Hertz Foundation. NV and DM are supported by a Simons Investigator Fellowship, DARPA grant W911NF2010021, and DOE grant DE-SC0022199. DM is supported by funding from the Office of Naval Research under award N00014-22-1-2377 and the National Science Foundation Grant under award CCF-2212841. SS is supported by a Susan Wojcicki and Dennis Troper Graduate Fellowship. BB is supported by a Google PhD fellowship. NV, BB, DM and CP are supported by funding from NSF grant DMS-2134157. CP is also supported by NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. 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. Compute was provided by the Harvard FASRC cluster and the Kempner Institute.
[51] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=1NvflqAdoom.
[52] Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186, 2023.
[53] Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1):2914, 2021.
[54] Zhengdao Chen, Grant Rotskoff, Joan Bruna, and Eric Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. Advances in Neural Information Processing Systems, 33:22217-22230, 2020.
A.1 MLPs
In Figure 9 we used a 3-layer MLP learning a Gegenbauer polynomial  $Q_{2}(\boldsymbol {\beta}\cdot \boldsymbol {x})$  in  $D = 5$  dimensions. Here  $\beta$  was a randomly chosen unit vector in  $\mathbb{R}^D$ . We implemented  $\mu \mathrm{P}$  parameterization by hand. The output layer of the network was rescaled by  $\alpha_0 / \sqrt{N}$ , consistent with  $\mu \mathrm{P}$ . We chose  $\alpha_0 = 1000$  to put us in the lazy regime. We set the learning rate to be  $5N / (1 + \alpha_0^2)$ . General arguments based on kernel scale indicate that the learning rate should be scaled as  $\alpha_0^{-2}$  at large  $\alpha$ .
In Figure 14 we used a 3-layer MLP learning a Gegenbauer polynomial  $Q_{2}(\boldsymbol{\beta} \cdot \boldsymbol{x})$  in  $D = 25$  dimensions. We set the learning rate to be nearly as high as possible before a loss explosion.
A.2 Vision
A.2.1 CIFAR-5m
All plots except Figure 6a and 6b: We trained with standard CIFAR data augmentation of random crop (RandomCrop(32, padding=4) in pytorch) and horizontal flip (RandomHorizontalFlip() in pytorch). As base network (for  $\mu P$ ) we used ResNet18 where BatchNorm was replaced with LayerNorm (to maintain the consistency of the neural network between train and test). We used the SGD optimizer with learning rate of .05 with cosine decay over 20000 steps, .9 momentum and batch size of 250.
For Figure 6a, we used the above setup, but with a learning rate of 0.01 and a much higher batch size of 2000, so as to replicate the edge of stability phenomenon [7] which only occurs at high batch sizes. For Figure 6b, we used a learning rate of 0.3 and batch size of 32, so as to show the behavior of high learning rate and small batch size on train loss.
A.2.2 CIFAR-10 Multiple Passes
In Figure 10, we show the dynamics and representational consistency of ResNets trained on CIFAR-10 for several epochs. The architecture is a ResNet-18 with base-shape width set at  $N = 64$  channels. The model is trained with SGD with learning rate 0.1 and cosine annealing schedule. The batch-size used is 128.
A.2.3 ImageNet
In all ImageNet experiments, we used a training subset of the ImageNet-1k dataset consisting of  $2^{20} = 1048576$  labeled images and a test subset consisting of 1024 labeled images. Both subsets were randomly sampled from the full ImageNet-1k training and validation datasets, respectively. To extend the duration in training in which the network remains in the online regime beyond one epoch, we heavily augmented the images in the training dataset using PyTorch's AutoAugment transform with the default policy, AutoAugmentPolicy.Imagenet.
We again used the ResNet-18 architecture with  $\mu \mathrm{P}$  parameterization relative to the ResNet-18 network with base-shape width  $N = 64$  channel [13]. All architectures and training procedures were implemented in Jax and used the auxiliary Flax and Optax packages, respectively.
Figures 2(b) and 7(b) were trained using the Adam optimizer with the following learning rate schedule: linear warm-up for 0.5 epochs from learning rate  $8 \times 10^{-5}$  to  $8 \times 10^{-3}$ , followed by cosine decay over 49.5 epochs to  $8 \times 10^{-5}$ .
A.3 Language
A.3.1 Wikitext-103 Language Modeling
For all Wikitext-103 tasks, we adopted the  $\mu \mathrm{P}$  transformer as defined in the  $\mu \mathrm{P}$  package [13]. In the plots shown in the main text, we used a depth-4 transformer, with  $d_{model}, d_k, d_v = N$  and  $d_{fn} = 4N$ . We performed a single pass through the train set in order to stay in the realistic online regime. We used a masked language modeling with sequence length  $S$  at varying input sequence lengths  $S$ . For Figure 1 d) we used the  $S \times S$  attention matrix of an  $S = 128$  transformer. In Figure 5 e) we used the attention matrix of an  $S = 35$  transformer. We chose this different length simply to
illustrate the consistent message across sequence lengths. We used a batch size of  $B = 32$  for all experiments. The residual stream was thus a tensor of shape  $(S, B, d_{model})$ .
We used the Adam optimizer with a learning rate of 0.0001. We also ran the same configuration with SGD and a learning rate of 0.5 and observed the same behavior. See section B for further plots and details.
For figure 3, we used the Wikitext-103 validation set in order to measure the evolution of the predictions on masked logits. In 3 f), we averaged the mean squared error from the widest transformer by using 100 test points.
A.3.2 C4 Language Modelling
Figure 1 (b) we trained with base network being a  $125\mathrm{m}$  parameter transformer model on 2.5 billion tokens using the Mosaic ML's LLM codebase (https://web.archive.org/web/20230519184343/https://github.com/mosaicml/examples/tree/main/examples/11m). See https://web.archive.org/web/20230519183813/https://github.com/mosaicml/examples/blob/main/examples/11m/yamls/mosaic_gpt/125m.yaml for the full hyperparameter details. We were limited by time and computational resources in our ability to explore further details of the C4 transformer model.
B Further Plots of Convergence
(a) Error Consistency
(b) Last Layer Preactivation Consistency
Figure 10: Training on CIFAR-10 in a ResNet-18 for multiple epochs generates dynamic preactivation densities and feature kernels which converge at realistic widths. (a) The test classification error curves for single models (solid) and ensembled (dashed) converge for realistic widths. (b) At initialization preactivation distributions in the last hidden layer of the CNN for a randomly chosen data point are Gaussian (as expected) and are very consistent across model widths  $N$ . To obtain histograms we train an ensemble of  $E = 8$  independently initialized networks concatenate activation patterns across members of the ensemble. After 20 epochs of training (models are around  $\sim 95\%$  accuracy), the preactivation distributions for the same data point have become non-Gaussian (consistent with infinite width theory) but are still remarkably consistent for large widths. (c) The final layer's feature kernel at initialization shows very little structure, but (d) after training networks of all widths converge to similar kernels. The plot in (d) compares ensemble averaged kernels with the  $N = 512$  ensembled kernel.
Figure 11: Convergence of layerwise representations in each layer (block)  $\ell$  of the ResNet-18 at large width  $N$  after training on CIFAR-5M.
In this section, we show additional figures illustrating convergence of network quantities across widths that we did not have space for in the main text.
B.1 Vision
A simple setting in which convergence properties are particularly clear and simple to study is for a ResNet learning CIFAR-10 and going over multiple passes of the dataset. In Figure 10 we plot a 20-epoch pass over CIFAR 10, and study the generalization error, initial and final preactivations in the last layer, and final layer kernels across widths. The training error begins to exhibit pathologies after sufficiently many epochs, related to the discussion in section 3.2.
Next in figure 11, we show a higher-resolution plot of the kernel Gram matrices across widths and across layers for the CIFAR-5M ResNet after a pass through the data. The larger resolution allows one to see that even the fine-grained details in the structure of the Gram matrix are consistent across widths.
B.2 Language
Next, in Figure 12, we create an analog of the language column of Figure 3, this time for  $\mu \mathrm{P}$  transformers of the same architecture and dataset but now optimized with vanilla SGD. The fact that wider transformers perform better still holds, and one can clearly see narrower networks approaching wider ones in their output logit values.
(a) Training Loss
(b) Correct logit value
(c) RMSE from widest
Figure 12: An analog of Figure 3 for  $\mu \mathrm{P}$  transformers trained with SGD. a) Training loss. It is interesting that in  $\mu$  parameterization the SGD optimized network is competitive with the Adam-optimized network. b) Value placed on the correct logit for a specific masked token. c) RMSE of correct logit value from the widest network.
C Defining  $\mu \mathbf{P}$  and SP (Standard Parameterization)
There are several detailed discussions about  $\mu \mathrm{P}$  vs SP scaling [12, 15, 9, 10, 14]. The aim of this section is to simply give an accessible and conceptual overview of their distinction, as well as a motivation for  $\mu \mathrm{P}$  from the perspective of keeping features moving in time even at infinite width.
There are several equivalent ways of parameterizing neural networks that give rise to the same dynamical effects, whether in  $\mu$ -parameterization or standard parameterization. We give the definitions in the case of a single-output feed-forward network and demonstrate that SP and  $\mu \mathrm{P}$  give rise to  $\Theta (N^{-1 / 2}),\Theta (1)$  feature movement at initialization, respectively.
Generalizations to other architectures (ResNets, Transformers) are straightforward. For a detailed discussion see [12] and also [14].
C.1 SP
We assume all hidden layers have equal width  $N$ . Let the input space have dimension  $D$ . Let  $\mu$  be the index of the training point in the dataset. At each layer  $\ell$ , the pre-activation  $h_{\mu}^{\ell + 1}$  in layer  $\ell + 1$  is given by
$$
\boldsymbol {h} _ {\mu} ^ {\ell + 1} = \frac {1}{\sqrt {N}} \boldsymbol {W} ^ {\ell} \cdot \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right), \tag {2}
$$
where  $\phi$  is an element-wise non-linearity, often taken to be the ReLU function. Here the  $N^{-1/2}$  out front allows  $h_{\mu}^{\ell+1}$  to be  $\Theta(1)$  at initialization as  $N \to \infty$  by the law of large numbers. The output of the network  $f_{\mu}$  is then given by:
$$
f _ {\mu} = \frac {\alpha}{\sqrt {N}} \boldsymbol {w} ^ {L} \cdot \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right). \tag {3}
$$
Here again the  $N^{-1/2}$  scaling again yields that  $f_{\mu}$  will be  $\Theta(1)$  as  $N \to \infty$ . In SP,  $\alpha$  is taken to be 1, but we will keep it explicit as it plays an important role in distinguishing the parameterizations. It is the laziness parameter identified in [31]. The change in the function is given by
$$
\frac {d f _ {\mu}}{d t} = - \eta \sum_ {\nu} K _ {\mu \nu} \ell^ {\prime} \left(f _ {\nu}, y _ {\nu}\right). \tag {4}
$$
Here  $K_{\mu \nu} = \nabla_{\theta}f_{\mu}\cdot \nabla_{\theta}f_{\nu}$  is the NTK gram matrix.  $y^\nu$  is the true label.  $\ell$  is the loss function (e.g. MSE or crossentropy) and  $\ell^{\prime}$  is its derivative with respect to the first argument. The NTK is easily seen to be order  $\alpha^2$  and  $\ell^\prime$  is order 1 at small  $\alpha$ . In order to have the change in the function be  $\Theta (1)$  we set  $\eta = \eta_0 / \alpha^2$ .
Using the chain rule, one can directly see that the pre-activations evolve as [10, 3, 32]
$$
\frac {d \boldsymbol {h} ^ {\ell}}{d t} \sim \eta \frac {\alpha}{\sqrt {N}} = \frac {\eta_ {0}}{\alpha \sqrt {N}}. \tag {5}
$$
Thus, at large  $N$  and  $\alpha = 1$  the pre-activations of this network evolve as  $\Theta (N^{-1 / 2})$ . Consequently, at infinite width the feature do not evolve and infinitely wide networks in standard parameterization become kernel machines with the static and initialization-independent infinite-width NTK.
In many machine learning libraries, the factors of  $1 / \sqrt{N}$  are not explicitly placed in front of each multiplication with the weight matrices. Rather, the weight matrices themselves are drawn from a distribution  $W^{\ell} \sim \mathcal{N}(0, \frac{1}{N} \mathbf{1})$ . Although this gives identical forward pass, this changes the  $N$  scaling of the gradients in the backward pass. As long as the learning rate is appropriately rescaled to account for this, the dynamics are equivalent to the SP parameterization discussed above.
C.2  $\mu \mathbf{P}$
One of the simplest ways to define the  $\mu$ -parameterization is to take  $\alpha = 1 / \sqrt{N}$ . This implies that we simply replace the final layer of the network by:
$$
f _ {\mu} = \frac {1}{N} \boldsymbol {w} ^ {L} \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right). \tag {6}
$$
As the prior analysis shows, in order to have  $df_{\mu} / dt$  be  $\Theta(1)$  at initialization, we take  $\eta = N\eta_0$ , so the learning rate in this definition scales extensively with  $N$ . In this setting, we now have that
$$
\frac {d \boldsymbol {h} ^ {\ell}}{d t} \sim \Theta (1). \tag {7}
$$