Feature-Learning Networks Are Consistent Across Widths At Realistic Scales

Introduction

Studies of large-scale language and vision models have shown that models with a larger number of parameters achieve better performance . 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 gives rise to a 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 . 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 . In this work, we focus on the maximal update parameterization ($\mu$P), or equivalently the mean field parameterization . 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?

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

For simple vision tasks such as CIFAR-5m , 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 5).

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 ) 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 , the emergence of outliers , and the emergence of induction heads 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 . 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 . 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$ for width $N$. A $1/N$ scaling of generalization error with width was empirically verified on many tasks . The effect of width on generalization in the feature-learning regime was empirically studied in 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 for linear models.

The authors of 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 . 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 . Extensions of this model to deeper networks were studied in . A theory of finite-width corrections to networks in this parameterization was studied in . A very general set of parameterization principles, termed $\mu$P, was introduced to give a well defined feature learning limit for a wide range of architectures including RNNs, CNNs, MLPs and transformers . 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$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.

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 , $N=d_{model}=hd_{k}=hd_{v}$ and $d_{ffn}=4N$. Here, $h$ is the number of heads, which we will keep 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).

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 . 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 , and more recent models such as GPT-3 have widths going up to 12288 , so this is also within the regime of realistic width.

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.

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

In Figure 6(a), 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 the edge-of-stability phenomenon identified in . We also show in Figure 6(b) 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$P networks. In Appendix 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$P leads to strikingly consistent loss curves, SP does not.

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 . This initialization-dependent variance can be mitigated by averaging the output logits of a sufficiently large ensemble of networks . 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.

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. ensembling 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 has focused mostly on studying the variance term. In Section 5 we study the bias from a theoretical perspective.

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 . 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 I. 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 as training progresses (Figure 8 (f)). Notably, we observe significant 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 for ensembled networks. We hypothesize that such counterexamples can only exist in the context of offline training.

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 $\text{Loss}=CN^{-\beta}+D$ where $C,\beta,D$ are fit to the data using the ‘scipy.optimize‘ function. The resulting fits are provided in Figure 9. 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 inifinite 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 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$ . This motivates novel theoretical descriptions of finite width mean field learning networks at late time which can capture task-dependent exponents.

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 .

Conclusion

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.

Acknowledgments and Disclosure of Funding

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.

References

Appendix A Experimental Details

In Figure 15 we used a 3-layer MLP learning a Gegenbauer polynomial $Q_{2}(\bm{\beta}\cdot\bm{x})$ in $D=25$ dimensions. We set the learning rate to be nearly as high as possible before a loss explosion.

A.2 Vision

All plots except Figure 6(a) and 6(b): 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 6(a), 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 which only occurs at high batch sizes. For Figure 6(b), 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 11, 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$P parameterization relative to the ResNet-18 network with base-shape width $N=64$ channel . 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

For all Wikitext-103 tasks, we adopted the $\mu$P transformer as defined in the $\mu$P package . In the plots shown in the main text, we used a depth-$4$ transformer, with $d_{model},d_{k},d_{v}=N$ and $d_{ffn}=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 125m 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/llm). See https://web.archive.org/web/20230519183813/https://github.com/mosaicml/examples/blob/main/examples/llm/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.

Appendix B Further Plots of Convergence

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.

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 11 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 12, 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 13, we create an analog of the language column of Figure 3, this time for $\mu$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.

Appendix C Defining μ𝜇\muP and SP (Standard Parameterization)

There are several detailed discussions about $\mu$P vs SP scaling . The aim of this section is to simply give an accessible and conceptual overview of their distinction, as well as a motivation for $\mu$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$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 and also .

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 . The change in the function is given by

Using the chain rule, one can directly see that the pre-activations evolve as

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.

C.2 μ𝜇\muP

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:

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

In , gives an equivalent definition of $\mu$P that gives rise to the same dynamics but keeps the learning rate to be $\Theta(1)$. We use this version of $\mu$P in the experiments that we run, simply because that is what is used in the package . Consequently, our learning rate does not need to be changed as width grows.

Appendix D Importance of Parameterization

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 14 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:

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$P train slightly faster.

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 .

Appendix E Further Studies of μ𝜇\muP versus SP

In figure 15, 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-indepdent 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 , 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.

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$P parameterized network in Figure 6(a). By contrast to that figure, we see significantly different dynamics and batch variation across widths. In Figure 16 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 6(a) b) . Further, at late times in Figure 16, 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 . We again stress that our observation is that not only are the final losses similar across widths in $\mu$P (as observed in ), but that the individual batch and large learning rate fluctuations agree across widths at early times in $\mu$P as well.

E.3 Language

Finally, we present a complementary set of figures to those in the right columns of Figures 3 and 13 for transformers of the same architecture on Wikitext-103 but in standard parameterization.

Appendix F Bias-Variance Decompositions over Initializations

In this section we explicitly define what we mean by initilization bias and initialization variance. Following , we consider the trained neural network function $f_{\theta^{*}}(\bm{x})$ to depend on the training set $\mathcal{D}$ (including inputs $\bm{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 ${(\bm{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 .

Defining the ensemble averaged netowrk predictor by

the expected squared bias of a neural network over initializations is:

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:

The variance of $E$ ensembles of a given network with independent initialization seeds $\theta_{e}$ is given by

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 the network ensembles decreases as $1/E$. In practice, we see that comparable decay rates are achieved.

Appendix G 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$P

G.2 Finite Width Effects

where $\mathcal{T}$ is the time-ordering operator. We now consider the leading correction to the average NTK around infinite width $\left<\bm{K}(t)\right>=\bm{K}_{\infty}(t)+\frac{1}{N}\bm{K}^{1}(t)+\Theta(N^{-2})$. With this correction, we see that the dynamics of errors $\bm{\Delta}$

The fact that the $\frac{1}{N}\bm{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 which tried to motivate a study of perturbation theory in logarithms of the transition matrix $\log\bm{T}(t)$ defined as

The solution to this can be used to construct the errors at a later time $\bm{\Delta}(t)=\exp\left(\bm{R}(t)\right)\bm{y}$.

Appendix H The sufficiency of small ensembles

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

Appendix I 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.

Appendix J 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 11 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 6(a) was just run for a few 100 steps of the training, so didn’t use much compute power.