SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nikhil Vyas $^{1*}$  Alexander Atanasov $^{2,3,4*}$  Blake Bordelon $^{1,3,4*}$
Depen Morwani $^{1,3}$  Sabarish Sainathan $^{1,3,4}$  Cengiz Pehlevan $^{1,3,4}$
$^{1}$ SEAS  $^{2}$ Department of Physics  $^{3}$ Kempner Institute  $^{4}$ Center for Brain Science
{nikhil,atanasov,blade_bordelon,dmorwani,
sabarish_sainathan,cpehlevan}@g.harvard.edu
We study the effect of width on the dynamics of feature-learning neural networks across a variety of architectures and datasets. Early in training, wide neural networks trained on online data have not only identical loss curves but also agree in their point-wise test predictions throughout training. For simple tasks such as CIFAR-5m this holds throughout training for networks of realistic widths. We also show that structural properties of the models, including internal representations, preactivation distributions, edge of stability phenomena, and large learning rate effects are consistent across large widths. This motivates the hypothesis that phenomena seen in realistic models can be captured by infinite-width, feature-learning limits. For harder tasks (such as ImageNet and language modeling), and later training times, finite-width deviations grow systematically. Two distinct effects cause these deviations across widths. First, the network output has an initialization-dependent variance scaling inversely with width, which can be removed by ensembling networks. We observe, however, that ensembles of narrower networks perform worse than a single wide network. We call this the bias of narrower width. We conclude with a spectral perspective on the origin of this finite-width bias.
Studies of large-scale language and vision models have shown that models with a larger number of parameters achieve better performance [1, 2]. Motivated by the success of large-scale models, several theories of deep learning have been developed, including large-width limits. Infinite width limits which arise in standard parameterization (SP) or neural tangent parameterization (NTP) considered in [3, 4] gives rise to a model with no feature learning (a kernel method). In this limit, the neural network loses the ability to adapt its internal features. Feature learning is crucial to explain deep learning's superior performance to kernels, the emergence of interpretable neurons such as edge-detecting CNN filters, transfer learning capabilities, and large learning rate effects such as edge of stability [5-8]. All of these effects are exhibited in modern large-scale networks.
Recently, several works have identified an alternative parameterization of neural networks that preserves feature-learning even at infinite width [9-14]. In this work, we focus on the maximal update parameterization  $(\mu \mathrm{P})$ , or equivalently the mean field parameterization [12, 15, 14]. The existence of infinite-width feature-learning limit, suggests that this parameterization is potentially more promising to explain deep learning phenomena than the previous limits. This motivates us to ask:
Question: Can realistic-width neural networks be accurately described by their infinite-width feature-learning limits?
(a) CIFAR-5m
(b) C4
(c) ImageNet
Figure 1: Consistency of large width behavior across tasks, architectures, observables. a) Loss curves for Resnets on CIFar-5M in  $\mu \mathrm{P}$  are nearly to identical at large widths (see also Figure 2). b) For GPT-2 on the C4 dataset [16] the loss curves agree at early times and deviate at late times, but wider networks agree for longer (see also Figure 2 and appendices for Wikitext-103) c) The values that ResNets put on the correct logit for ImageNet appear to converge as the width grows (see also Figure 3). d) The attention matrices for transformers on Wikitext-103 become nearly identical as width increases (for quantitative metrics see Figure 5.)
(d) Wikitext-103
We attempt to answer this question by training networks of varying widths on vision and language tasks for realistic datasets and architectures. Concretely, we focus on on the online setting, where data is not repeated during SGD, and track the following quantities across widths:
- The losses throughout training.
- The predictions of the networks on individual points throughout training.
- The learned representations, summarized by the feature kernels preactivation distributions and, for transformers, their attention matrices.
- Dynamical phenomena such as the edge of stability governing the top Hessian eigenvalues, as well as large learning rate and small batch size effects on the loss.
On each of these metrics, we show that sufficiently wide neural networks converge to consistent behavior across widths. In Figure 1, we show loss curves, logit predictions, and attention matrices approach consistent behavior as width is increased across several architectures and datasets. We further observe that the widths that achieve this consistent behavior are within the range of those used in practice. We use large-width consistency as a proxy for achieving the limiting infinite-width behavior. We stress that this observed consistency is a property of networks in mean field/ $\mu$ P parameterization but is not present in other parameterizations which also give an infinite width limit like NTK parameterization (See Appendix D for a comparison).
We say that a network property is consistent if, beyond some width, its values all lie within some small interval with high probability. We measure consistency by showing that a quantity's deviations between successive widths decrease as the widths are increased, and that its value for narrower networks systematically approaches its value for the largest trained network.
Our results show the following:
- For simple vision tasks such as CIFAR-5m [17], ResNets with practical widths achieve near consistent loss curves across widths (Section 2).
- Beyond the loss curves, the individual predictions of the networks agree pointwise. That is, the logits agree on test points throughout the training process. We further show that internal representations as measured by distributions of neuron preactivations and feature kernels in various layers are consistent across widths (Section 2).
- For harder tasks such as ImageNet and language modeling, loss curves are consistent across widths early in training. As training progresses, loss curves for narrow networks deviate smoothly from the loss curves of wider networks. The effective width required to reach infinite-width behavior thus increases with training time. Conversely, as network size grows we approximate the infinite width network for a larger number of training steps (Section 2).
- Finite-width neural networks have variance in the learned function due to initialization seed. This variance depends inversely on the width. We study ensembles of networks over different initializations to remove this noise. Further, by training ensembles of networks, we can perform
a bias-variance decomposition over initializations (c.f. Appendix F for details and definitions) to analyze the effects of finite width. We find that finite-width bias plays an important role. Equivalently, ensembling narrow networks does not yield infinite-width behavior (Section 3).
- In the setting of offline learning, at late times one can over-fit the training set. We observe that this leads to larger gaps in network behavior across widths, and can break the trend that wider networks perform better (Section 3).
- We develop a spectral perspective on the origin of the finite-width bias by analyzing it in a simple setting of a lazy network learning a simple task. We then apply this perspective to a CNN trained on CIFAR-5m (Section 4).
The consistency across large widths strongly suggests that the dynamics and predictions of realistic-scale networks can be effectively captured by their infinite-width feature learning limits. For realistic tasks, as the width is increased, a larger interval of training can be characterized by this infinite-width limit. Most importantly, even though quantitative agreement across widths slowly breaks with more and more training, we observe that wider networks perform better (as in [15]) and preserve qualitative aspects of the learned features (such as hidden layer kernels and attention matrices) and dynamical phenomena such as edge of stability. This suggests that infinite width feature-learning networks are good models to study deep learning.
Our results have implications for interpretability, as the agreement of internal representations suggest that many other phenomena, such as transfer learning with linear probes or fine-tuning, in-context learning [18, 19], the emergence of outliers [20], and the emergence of induction heads [21] may be understood from the perspective of infinite-width feature learning networks.
We plan to have our code made freely available on github to ensure the reproducibility of these results.
Empirically, the scaling of relevant quantities with width in the standard or neural-tangent parameterizations was thoroughly studied in [22]. In the latter parameterization, sufficiently wide networks give a kernel method with the infinite-width NTK. Several papers have shown that in practice the NTK limit insufficiently characterizes realistic deep neural networks [5, 23, 6]. Attempts to capture feature learning and predictor variance from perturbative series around infinite-width dynamics show that finite-width variance and kernel adaptation scale as  $1 / N$  [24-26] for width  $N$ . A  $1 / N$  scaling of generalization error with width was empirically verified on many tasks [27, 28]. The effect of width on generalization in the feature-learning regime was empirically studied in [29] in the relatively limited setting of multi-layer perceptrons (MLPs) on polynomial tasks. There, the variance of the finite-width NTK at the end of training adversely affected generalization. Bias-variance decompositions over dataset, label noise, and initialization parameters were studied in [30] for linear models.
The authors of [31] identified that altering the output scale  $\alpha$  of any network could increase or decrease feature learning in a neural network. Large values of  $\alpha$  correspond to the "lazy limit" where the network's features don't evolve. A follow up study noticed that rescaling the output by  $\alpha = \alpha_0 / \sqrt{N}$  for width  $N$  networks gave consistent behavior of feature learning and losses in small scale experiments [10]. Several works have studied this regime of training in the two-layer limit, known as "mean field" parameterization, where features are still learned even at infinite width [32, 9, 33, 34]. Extensions of this model to deeper networks were studied in [35-38, 12, 14]. A theory of finite-width corrections to networks in this parameterization was studied in [39]. A very general set of parameterization principles, termed  $\mu \mathrm{P}$ , was introduced to give a well defined feature learning limit for a wide range of architectures including RNNs, CNNs, MLPs and transformers [12]. [15] demonstrated that this parameterization nearly fixes optimal hyperparameters across network widths, allowing for hyperparameter transfer from small to large widths. This work also empirically noted that wider networks always outperformed narrower networks in this parameterization.
Our paper focuses on networks in  $\mu \mathrm{P}$  and attempts to study the consistency of many relevant network properties across widths. We perform a fine-grained analyses of more realistic models throughout the dynamics of training. To the best of our knowledge, this is the first such paper to study the consistency of network outputs, internal representations, and dynamics across widths.
(b) Imagenet
(c) Wikitext-103
Figure 2: In the online learning setting, train loss improves as width grows. For sufficiently wide networks, the training lost is consistent across widths. For Cifar-5m this consistency is observed over all of training. For harder tasks like Imagenet and Wikitext-103, networks of different widths agree up until a width-dependent time-step where narrower networks begin performing worse.
(a) CIFAR-5m
(d) CIFAR-5m
Figure 3: The output logits on a fixed test point displays stable behavior at large enough widths. a) Value of network on correct class logit over time as width is varied for CIFAR-5m. Colored errorbars represent one standard deviation. b) Same plot for Imagenet for a fixed image in the test set c) Same plot for Wikitext-103 for a fixed masked token. Across the board the widest networks behave similarly. Next, we use the widest network as a proxy for the infinite-width limit, and compare the logit predictions of narrower networks against that. d) For CIFAR-5m, the relative root-mean-squared error over the test set of the distance to the value that the widest network puts on the correct logit. e) The same for Imagenet. f) The same for Wikitext-103. We see a striking regularity of networks converging to the widest one as the width grows. In Appendix B, we also compare networks of successive widths and show the difference shrinks.
(b) Imagenet
(e) Imagenet
(f) Wikitext-103
2 Consistency of large-width behavior in online learning
We focus on studying the effect of width in the setting of neural networks learning a task in the online setting. Online learning is representative of many modern settings of deep learning, and as will be shown in Section 3, obviates consideration of memorization and over-fitting in offline learning that can lead to large differences in networks across widths.
In what follows, the variable  $N$  will denote the width of a given network. For vision tasks, this will correspond to the number of channels in each layer. For transformers, in the notation of [40],  $N = d_{model} = hd_k = hd_v$  and  $d_{ffn} = 4N$ . Here,  $h$  is the number of heads, which we will keep
(b) Imagenet
Figure 4: Analog of the last row of Figure 3 but comparing networks of successive widths rather than comparing all networks to the widest. Again, we see that as the network width grows, the difference between successive networks shrinks.
fixed.  $d_{model}$  is the embedding dimension of the tokens as well as the dimension of the residual stream.  $d_{k}$  is the dimension over which the dot products in the attention are calculated and  $d_{v}$  is the dimension of the values in the attention layers.  $d_{ffn}$  is the hidden width of the feedforward networks (FFN).
Convergence of loss curves We begin by showing (Fig. 2) that the loss curves for sufficiently wide networks on a given task achieve consistent behavior across widths. Throughout the paper we measure train loss in terms of crossentropy. For all tasks, at early times large widths agree, but for more complicated tasks such as ImageNet or Wikitext-103, learning curves of narrower network deviate from those of wider ones.
The width beyond which networks emulate infinite-width behavior depends on the complexity of the task. For more difficult tasks, larger widths are required for the loss curves to converge. For simple tasks such as CIFAR-5m we find that widths as narrow as 128 are essentially consistent with infinite width-behavior for an entire pass through the 5 million image dataset. For ImageNet, widths near 512 are close to consistent for four passes through the dataset with heavy data augmentation. These widths are well within the range of those practically for images [41, 42]. For transformers going through a single full pass of Wikitext-103, widths on the order of 4000 are required. Early transformer models certainly had hidden widths of order 4k [43], and more recent models such as GPT-3 have widths going up to 12288 [18], so this is also within the regime of realistic width.
Pointwise convergence of predictions Beyond the convergence of the training loss curves, we observe that the logits of a network on a fixed test point become consistent as width grows. This test point can be an image in the test set or a masked token in the validation set. In plots a), b), and c) of Figure 3, we show that for a specific held-out test point, the value of the network on the correct logit becomes consistent as the width grows. In d), e), and f) we plot the root mean squared distance to the widest networks logits over the test set. We further study the difference between successive widths in Figure 4.
Convergence of representations In addition to loss and prediction dynamics, we also examine whether learned representations in these models are consistent across widths. Mean field theories of neural network dynamics predict that sufficiently wide networks should have identical kernels (and attention matrices for transformers) and that all neurons in a layer behave as independent draws from an initialization-independent single-site distribution [9, 12, 14, 46, 47]. To test whether realistic finite-width feature learning networks are accurately captured by this limit, in Figure 5, we analyze the feature kernels and preactivation distributions before and after training as well as the attention matrices in transformer models trained on Wikitext-103. We see qualitative consistency in the plots of kernels and attention matrices in b) and c) which can be made quantitatively precise by plotting the distance to the widest networks and showing systematic convergence in c) and f).
Convergence of dynamical phenomena In Figure 6a, we show that the sharpness, defined as the top eigenvalue of the loss Hessian, grows steadily to a final value that it then fluctuates around. This is a small-batch analog of the edge-of-stability phenomenon identified in [7]. We also show in Figure 6b that on CIFAR-5m task, at early times, the individual variations due to batch noise and large learning rate effects can be consistently captured across widths for  $\mu \mathrm{P}$  networks. In Appendix
(a) Preactivations CIFAR-5M
(b) Last Layer Kernels
(c) Convergence of Kernels
(d) FFN Preactivations Wikitext
(e) Attention Matrices
(f) Convergence of Attention
Figure 5: Learned features are consistent across a large range of widths in realistic tasks. (a) The distribution (over neurons) of preactivation values  $h$  in the final block of  $E = 8$  ResNet18 networks trained on CIFAR-5M. At initialization, the densities are all well approximated by the Gaussian with matching mean and variance (dashed black). After feature learning, the density has shifted and become non-Gaussian (poor match with dashed black), yet is still strikingly consistent across widths. (b) Average (over random init) feature kernels are also consistent across widths. (c) The centered kernel alignment CKA [44, 45] of the width  $N$  and width 512 kernels increases towards 1.0 as  $N$  increases. The  $1 / \sqrt{N}$  and  $1 / N$  trends are plotted for reference. (d) The preactivation histogram for a transformer on Wikitext-103. At initialization the Gaussian of best fit is the standard normal. After training the histograms are still quite Gaussian, with different moments. (e) A variant of Figure 1 (d) at a smaller sequence length. Attention matrices are consistent at large widths. f) Both FFN kernels and attention matrices converge as width grows. The  $1 / N$  and  $1 / \sqrt{N}$  trends are plotted for reference.
(a) Edge of Stability
Figure 6: Convergence of dynamical phenomena across width for CIFAR-5m
(b) Small Batch size effects
E, we further demonstrate sharp agreement of large learning rate and small batch size phenomena for MLPs learning a simple task. There, we show that while  $\mu \mathrm{P}$  leads to strikingly consistent loss curves, SP does not.
3 Deviations from large-width behavior
The consistency observed in Section 2 may break later during training in either the online or offline settings. In the online setting, deviations owing to narrow width compound over time and lead to two sources of error relative to the infinite width limit which we describe in 3.1. In the offline setting, where data is recycled several times, networks over-fit the training data, which can lead to larger gaps between widths and can break the trend that wider networks perform better.
Finite-width effects introduce an initialization dependence to the network, leading to additional variance in the learned function and hindering generalization [27-29]. This initialization-dependent variance can be mitigated by averaging the output logits of a sufficiently large ensemble of networks [48]. Using the bias-variance decomposition terminology, we refer to the discrepancy in performance between an ensembled network and the expected performance of a single network the variance, and the gap between an ensembled network and the behavior of infinite-width network as the bias of narrower width. We elaborate thoroughly on what we mean by this decomposition in Appendix F. By definition, the expected difference in loss between a single finite-width network and an infinite-width network is the sum of the bias and the variance. Below, we investigate the behavior of bias and variance in networks across various vision and language tasks.
3.1 Online training
(b) Imagenet
Figure 7: Loss curves and their ensembles in the online setting. Ensembling reduces the training loss, but a large ensemble of narrow networks do not achieve the performance of a single wider network. Errorbars representing standard deviation over random init are not visible.
Figure 7 shows that at large widths, both single networks and ensembles of networks achieve comparable error. In this regime, all the networks are consistent and increasing the width has a very marginal effect, as does ensembling. At narrower widths, variance is nontrivial (i.e. assembling helps) but bias is much larger than variance. Single wide networks outperform ensembles of narrower networks. By comparing a) with b) and c) of Figure 7, we see that harder tasks induce larger bias gaps. Prior theoretical work [28, 29] has focused mostly on studying the variance term. In Section 4 we study the bias from a theoretical perspective.
3.2 Offline Training
In offline learning, which refers to multi-epoch training, we encounter several unexpected phenomena that challenge the width consistency observed in the previous section, even at large widths. To compare offline learning with online learning, we utilize CIFAR-200k, a 200k sized random subset of CIFAR-5m. Previous studies have demonstrated that label noise contributes to an increase in overfitting [49]. In order to investigate how width consistency changes with overfitting and double descent, we conduct experiments on a noisy label version of CIFAR-50k (50k sample from CIFAR-5m), where  $50\%$  of the labels are noisy. Additional ImageNet experiments are presented in Appendix J. As offline training achieves near-zero error, we need to compare very small quantities. To accomplish this, we will plot and compare all quantities on a logarithmic scale. The following phenomena are observed:
- Single network performance on the training set does not converge with width, even at high widths (Figure 8 (a)). In other words, the combined bias and variance does not reach zero, even with substantial widths. This is in contrast to the online runs.
- Ensembling (Figure 8 (b)) reveals that both bias and variance terms individually fail to reach zero, even at high widths.
- Regarding test performance, both bias and variance tend to zero as width increases, demonstrating an instance of benign overfitting (Figure 8 (d) and (e)).
- When working with the noisy label version of CIFAR-50k, we observe clear overfitting and stepwise double descent [49] as training progresses (Figure 8 (f)). Notably, we observe significant
(d) CIFAR-200k, offline vs. online.
(e) CIFAR-200k, offline ensembled