Feature-Learning Networks Are Consistent Across Widths At Realistic Scales
Authors: Nikhil Vyas
Alexander Atanasov
Blake Bordelon
Depen Morwani
Sabarish Sainathan
Cengiz Pehlevan
Venue: NeurIPS 2023 poster (2023)
Keywords: mean-field
muP
feature learning
infinite width
deep ensembles
TL;DR: Realistic networks achieve infinite-width feature-learning behavior throughout training for easy tasks and for early training time in hard tasks.

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

Feature-Learning Networks Are Consistent Across Widths At Realistic Scales
Nikhil Vyas
Alexander Atanasov233.4
Blake Bordelon 1,3,4
Depen Morwani
Sabarish Sainathan 133,1
Cengiz Pehlevan
SEAS
Department Physics SKempner Institute
*Center for Brain Science
Harvard University
(nikhil,atanaso) atanasov, blake aellak_bbrddeee bordelon, dmorwani,
sabarish sainathan, cpehlevan harvard.edu
Abstract
We study the effect of width on the dynamics of feature-learning neural networks across variety of architectures and datasets. Early in training, wide neural net- works 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, pre- activation 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, Crature-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 detialization-deeeddee variance scaling inversely with width, which can be removed by ensembling net- works. 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 this finite-width bias.
Introduction
Studies of large-scale language and vision models have shown that models with larger number of parameters achieve better performance 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 B 4 gives rise to model with 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 edge- detecting CNN filters, transfer learning capabilities, and large learning rate effects such as edge of All of these effects exhibited in modern large-scale networks.
stability 55-
Recently, several works have identified an alternative parameterization neural networks that preserves fcature-learning even infinite width [9-14]. In this work, focus the maximal update parametcrization (μPP,, or cquivalently the mean field parameterization [22 13 14. The existence of
infinite-widtt feature-learnin 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?
TTheee authors contributed equally to this work.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

N- N24124
254
N- $12
1624
(a) CIFAR-5m
(b) C4
(c) ImageNet
(d) Wikitext-103
Figure 1: Consistency large width behavior across tasks, architectures, observables. a) Loss curves for Resnets on Cifar-5M in μP nearly to identical 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-1033 103) e) 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 Wikitext-100 become nearly identical width increases (for quantitative metrics see Figure
We attempt answer this question by training networks of varying widths 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 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 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 metries, we show that sufficiently wide neural networks converge to consistent behavior across widths. In Figure 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-widtt consistency as a proxy for achieving the limiting infinite-wiet behavior. We stress that this observed consistency is property of networks in mean field/μP parameterization but not present other parameterizations which also give an infinite width limit like NTK parameterization (See Appendix [ for comparison).
We say that network property consistent if, beyond some width, its values all lie within some small interval with high probability. We measure consistency by showing that quantity's 's deviations between successive widths decrease 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 ImageNet and language modeling, loss curves 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- wiith behavior thus increases with training time. Conversely, as network size grows we approximate the infinite width network for 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 variance decomposition over initializations (c.f. Appendix For details and definitions) analyze the effects of finite width. We find that inite-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 spectral perspective on the origin of the finite-width bias by analyzing it in 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, the width is increased, 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 (T5) 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 [1], the emergence of outliers [20], and the emergence of induction heads 21 may be understood from the perspective of widhh feature learning networks.
We plan to have our code made freely available on github to ensure the reproducibility of these results.
1.1 Related Works
Empirically, the scaling of relevant quantities with width in the standard or neural-tangent parameteri- zations 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 (. Attempts to capture feature learning and predictor variance from perturbative series around infinite-width dynamics show that finite-width width variance and kernel adaptation scale as 1/N (24 26) for width N. 1/N scaling of gen-
eralization error with width was empirically verified on many tasks [27 28 The effect of width on generalization in the feature-learnnng regime was empirically studied in 29 in the relatively limited setting of multi-layer perceptrons (MLPs) polynomial tasks. There, the variance of the finite- NTK at the end of training adversely affected generalization. dataset, label noise, and initialization parameters were studied in [30 for linear models.
decompositions over
The authors of 3T identified that altering the output scale a of any network could increase or decrease feature learning in neural network. Large values of a correspond to the "lazy limit" where the network's features don't evolve. follow up study noticed that rescaling the output by = ou/ //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
Extensions of this model to deeper networks were studied in [5-388 38 12 14] A
theory of wiith corrections to networks in this parameterization was studied in [39]. A very general set of parameterization principles, termed jP, was introduced to give well defined feature learning limit for a wide range of architectures including RNNs, CNNs, MLPs and transformers [2]. [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 jPP and attempts to study the consistency of many relevant network properties across widths. We perform fine-grained analyses of more realistic models throughout the dynamics of training. To the best of our knowledge, this the first such paper to study the consistency of network outputs, internal representations, and dynamics across widths.
3

---

Red32
20
224
387
16'
337
15
10*
(a) CIFAR-5m
(b) Imagenet
(c) Wikitext-103
In the online learning setting, train loss improves as width grows. For sufficiently wide networks, the training lost is consistent actoss 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 width-dependent time-step where narrower networks begin performing worse.
Figure 2:
B--24
323
M-120
M--25
3024
4-512
5000
10000
15000
20000
Tain
Smagees 8001
Sean
(a) CIFAR-5m
Imagenet
(c) Wikitext- 103
No.12
328
0.5
256
0.2
1024
0.20
5000
10000
15000
20000
04
2.5
Tel Supa
2nagga0 2000
(d) CIFAR-5m
Imagenet
(I) Wikitext- 103
Figure 3: The output logits on fixed test point diplays stable behavior at large enough widths. a) Value of network correct class logit over time as width is varied for CIFAR-5m. Colored errorbars represent one standard deviation. b) Same plot for Imagenet for fixed image in the test set c) Same plot for Wikitext-103 for fixed masked token. Across the board the widest networks behave similarly. Next, we use the widest network as a proxy for the infinite-win width limit, and compare the logit predictions of narrower networks against that. d) For CIFAR-5m, the relative root-mean -uuunnd етгог 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- converging to the widest one as the width grows. In Appendix B we also compare networks of successive widths and show the the difference shrinks.
We see a striking regularity of networks
2 Consistency of large-width behavior in online learning
We focus studying the effeet of width in the setting of neural networks learning 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 wert-fitting in offline learning that can lead to large differences networks across widths.
In what follows, the variable N will denote the width of given network. For vision tasks, this will correspond to the number of channels in each layer. For transformers, the notation of 40),
N dmode! hdd. hd, and djfn 4N. Here, h is the number of heads, which will keep

---

132. 56
K42221
190)
028 25
11.00
12%% 5221
C32 3300
Teth Stpps
365
TokensS S499
363
(a) CIFAR-5m
(b) Imagenet
(c) Wikitext-
Figure 4: Analog of the last of Figure but comparing networks of successive widths rather than comparing all networks to the widest. Agan, we see that as the network width grows, the difference between successive networks shrinks.
fixed. dmodee is the embedding dimension of the tokens well as the dimension of the residual stream. dy is the dimension over which the dot products in the attention are calculated and d, is the dimension of the values in the attention layers. d 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-108 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 required for the loss curves converge. For simple tasks such as CIFAR-5m we find that widths as narrow as 128 are essentially consistent with infinite width-behavior for entire pass through the 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 [1 421 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 (3), and more recent models such as GPT-3 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 ob- serve that the logits of network on a fixed point become consistent width grows. This test point can be an image in the test set 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 net- works logits over the test We further study the difference between successive widths in Figure
Convergence of representations In addition loss and prediction dynamics, also examine whether Iearned 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 layer behave as independent draws from an ietaaiizaaionnnndeeedddee finite-widhh feature learning networks are accurately captured by this limit, in Figure [S] we analyze the feature kernels and preactivation distributions before and after training as well as the attention matrices in transformer models trained on Wikitext- 1033. We see qualitative consistency in the plots of kernels and attention matrices in b) and e) which can be made quantitatively precise by plotting the distance to the widest networks and showing systematic convergence in c) and f).
single-site distribution 9 12
47 To test whether realistic
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 small-bated analog of the the edge-of-stablle Figure (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 pP networks. In Appendix
phenomenon identified in [7]. We also show in
5

---

Initial
Final
1-CKAlGf..
2009
24
10-l
(a) Prcactivations CIFAR-5M
(b) Last Layer Kernels
(c) Convergence of Kernels
12-
Ni 1024
N: 2016
N 6096
Artwwiin
Acthorion
2040
(d) FFN Preactivations Wikitext
(e) Attention Matrices
(I) Convergence Attention
Figure 5: Learned features consistent across large range of widths in realistic tasks. (a) The distribution (over neurons) of preactivation values in the final block of trained 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-Gate 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 increases. The 1/VN and 1/N trends plotted for reference. (d) The preactivation histogram for transformer on Wikitext- 103. At initialization the Gaussian of best fit is the standard normal. After training the histograms still quite Gaussian, with different moments. (e) variant of Figure [ (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 I/N and L/v/ trends are plotted for reference.
ResNet18 networks
44 45 of the width N and width 512 kernels increases towards 1.0 N
MWM MO
120
2.
Trar Strpe
(a) Edge Stability
(b) Small Batch size effects
Figure 6: Convergence of dynamical phenomena across width for CIFAR-5m
we further demonstrate sharp agreement of large learning rate and small batch size phenomena for MLPs learning simple task. There, show that while μ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 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 This matialization-deeedder 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 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 [ By definition, the expected difference in loss between single finite-width network and an infinite- 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
rigl Satkort
4004
32P
10
10'
38'
10
Tain
(a) CIFAR- 55m
(b) Imagenet
(c) Wikitext- 103
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 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. this regime, all the networks are consistent and increasing the width has a very marginal effect, does ensembling. At narrower widths, variance is nontrivial (i.e. ensembling helps) but bias much larger than variance. Single wide networks outperform ensembles of narrower networks. By comparing a) with b) and e) of Figure 7 we that harder tasks induce larger bias gaps. Prior theoretical work [28] 22
29 has focused mostly on studying the variance term. In Section study the bias from theoretical perspective.
3.2 Offline Training
In offline learning, which refers to multi-epoch training, 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, 200k sized random subset of CIFAR- 5m. Previous studies have demonstrated that label noise contributes to an increase in overfitting 44 order to investigate how width consistency changes with overfitting and double noisy label version of CIFAR-50K (50k sample from CIFAR-
descent, we conduct experiments 5m), where 50% of the labels are noisy. Additional ImageNet experiments are presented in Appendix[] As offline training achieves near-zero error, 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 (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 (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 (c)).
When working with the noisy label version of CIFAR-50K we observe clear overfitting and stepwise double descent [49 as training progresses (Figure 8 (D). Notably, we observe significant
7

---

130
Taann CHO0
5-2 Office
0.00
Se5an offine
an
100
100
1000
C0000
Stops
(a) CIFAR-200K; offline vs. online
(b) CIFAR-200K, offline ensembled
(c) Noisy CIFAR-5OK train eTRor
onera
1-130 Offire Extrbia
Eftire
64532 Offine
Trann 5t:p5
Train Stops
50000
1000
2000
40:00
8006
(d) CIFAR-200K, offline vs. online. (c) CIFAR-2000 offfine ensembled
(f) Noisy CIFAR 50k test error
Top Row: Effects of offline training on train metrics. In (a) and (b) 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 for the largest widths 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- We see that both bias and variance for test error are near O at large widths. (f) We train on noisy CIFAR-50K and observe that "wider is better" violated for ensembled networks.
Figure 8:
could try. (c) we do multi-epoch training
deviations in width for single network performance, indicating that the benign overfitting observed Figure (d) and (e) is dataset-dependent. Furthermore, variance is found be much larger than in the non-noisy experiments.
Surprisingly, we discover (Figure S (f) that some ensembled narrower width networks out- perform ensembled wider networks. This presents 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 first in the simple setting of an MLP fitting polynomial in the lazy limit. Here, all the dynamics are well-captured by the finite-wide empirical neural tangent kernel (cNTK). By studying the spectral properties of this kernel across widths, 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 52).
We consider a model of online learning where large batch of data from the population dis- tribution p(x:) is sampled each step. This leads to approximate gradient flow dynamics di - oo (z 0) {()))² (Appendix Hp. To analyze this equation, we choose a fixed orthonormal basis (x)) for the space T.2 T.2(RR p((r)dda)) of square-integrable functions on in- put space. The function fæ) residual error A(x) y(x) f(x). and the empirical NTK
8

---

((x,v', x',t) .t) can be expressed in this basis as f(z,t) t) Σx fo (2), A(z,t) = Σx (x)
and K (x, Exe (xe(t) (2) respectively. respeccivvl Their training evolution is given by:
d
di [x, Ee'n ())((a')
E (Ovi)t)
(I)
The statisties of the dynamical NTK matrix Kg(() summarizes the statistics of the error dynamics Axx t) at any level of feature learning. At infinite width, Kxe(t) (t) is deterministic, while at finite width, it receives (NN=1) mean displacement and @ (N- fluctuation around its mean 24 25 391. We consider approximating the dynamics of the ensembled predictor by d. (fa(t)),, s Ee (t)) (t)(.)). Here, (-) denotes averages over initializations. This expression neglects the contribution from Cov(K Kov( &e). We show that this approximation is accurate in depth-3 MLPs trained on Gegenbauer polynomial regression tasks in Figure(aa). For more details see Appendix [
di
In the lazy limit, the kernel is static and we choose t/e to diagonalize (Ker) = ògeAR- This yields the loss dynamics C(t) = En (x))) e - We can therefore quantify aligoment of eigen- functions to task with the cumulative power distribution C(k) Etek (v()))))(())) (v(())))) which is the proportion of the task that is explained by the first &s eigenvectors (53]. If C(k) rises rapidly with & 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 A which set the timescales could be width derrhddeeeddent. Second, the eigenfunctions 4/ (a) that diagonalize (K) can change with width. In Figures b) and c) we show that the dominant effect is the latter. Finite-widtt 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 slower to learn.
To test whether these findings continue 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 (Ka x²)) for output channels e, c' and input images E We then compute the kernel gradient flow corresponding to MSE training on the true =
target function for CIFAR-5M d (x) Ee Ex (Ke. xe)) A (x') from initial condition given by the one-hot target labels (x(=)):=00 (z(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 MIP experiment, the spectra are very consistent across widths, but the cigenfunction 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.
5 Conclusion
We have demonstrated 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-widih 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 compute fronticr. We leave further inquiry into such limits for future work.

---

Lasser br.art MLP end, beich 252
4. MLP. 201
1-04 MLD. Je2
M-105
34'
HH-23
16
23.4
341
3 11-
10
321
13
16-1
Exponvettr R0C4
MLP Online training losses
(b) MLP Kernel Spectra
(c) MLP unexplained task fraction
(3) Task Abgumeni
Kernel How Final NX1
(3 Spec.ra
33
12"
0.2
161
15"
237
17
14'
33a
(f) After-Kernel unexplained task fraction
(d) CIFAR-5M Kfter-Kemnel
(c) After-Kernel Spectra
Figure 9: Spectral properties of the NTK can account for bias gaps across widths. (a) Depth 3 MLPs in the lazy limit (7 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 iigendeccomposttion of the task has greater weight along higher spectral modes, consequently leading to slow-down in training. These results hold across dimensions, batch sizes, task complexity, and architectures. Strong feature learning can reduce this effect. See Appendix (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 kemels noticeably differ across widths, evidenced by the cumulative power distribution C(k) which accounts for the gap in theoretical loss curves under kernel flow.
Acknowledgments and Disclosure of Funding
We thank Boaz Barak, Jeremy Cohen, Alex Damian, Nikhil Ghosh, Gal Kaplun, Eric Michaud, Jamie Simon and Jacob Lavatone-Veth for helpful discussions throughout this project. We also thank Jacob Zavatone- Vetth 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 W91 INF20 MONN2000021.000 annd DOE grant DE-SC0022 199. DM is supported by funding from the Office of Naval Research under award N00014-22-1-233 award CCF-22118411. SS is supported by Susan Wojcicki and Dennis Troper Graduate Fellowship. BB is supported by Google PhD fellowship. NV, BB, DM and CP are supported by funding from NSF grant 2M3-2134157. CP is also supported by NSF CAREER Award IIS 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.
and the National Science Foundation Grant under
10

---

References
[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, arXiv preprint arXiv:2001.00 003361, 2020.
[2] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[3] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[4] Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[5] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Anmual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12. 2020, virtual, 2020. URL Abstrattttttttttttttc..c html
neurips cc htcppapprr/222//ee/a/
[6] Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the ntk for understanding generalization in deep learning. arXiv preprint arXiv:2206. 10012, 2022.
[7] Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gra- dient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.0000 00065, 2021.
[8] Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: Circuits. Distill, 2020. doi: 10.23915/ distill.00024. https://dsiillppb2220cirrrriri pub/2020/circuiss.
[9] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018.
[10] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 4020(11):111) 13301, 2020.
[11] Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.
[12] Greg Yang and Edward J Hu. Tensor programs Feature learning in infinite-width neural networks. In International Conference Machine Learning, pages 1 11727-11737. PMLR. 2021.
[13] Greg Yang and Edward Hu. Maximal Update Parametrization (nP) and Hyperparameter Transfer (uTrransfer). httpss//shhaa..cco//oooooorroogp 2022.
[14] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. arXiv preprint arXiv:2205.006533 2022.
[15] Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL hhtpps //openreview (ttorruunnniiddd622uu2222A2

---

[16] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555) 2020.
[17] Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. In International Conference on Learning Representations, 2021.
[18] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-19011, 2020.
[19] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud. Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0766 076822 2022.
[20] Ziyang Luo, Artur Kulmizev, and Xiaoxi Mao. Positional artefacts propagate through masked language model embeddings. arXiv preprint arXiv:2011.04392 2020.
[21] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[22] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. Advances in Neural Information Processing Systems, 15115166-15122 2020.
[23] Guillermo Ortiz-Jiménez, Seyed-Mohsen Mohsen Moosavi-Dezfooli, and Pascal Frossard. What can linearized neural networks actually say about generalization? Advances in Neural Information Processing Systems, 34:8998-9010, 2021.
[24] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint arXiv: (1999099998) 2019.
[25] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[26] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. Cambridge University Press Cambridge, MA, USA, 2022.
[27] Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d' Ascoli, Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2020(2).033001, 2020.
[28] Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701, 2021.
[29] Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The on- set of seriance-limited behavior for networks in the lazy and rich regimes. arXiv preprint wrXi::2212.12147. 2022.
[30] Ben Adlam and Jeffrey Pennington. Understanding double descent requires a fine-grained bias- variance decomposition. Advances in neural information processing systems, 33:1 1022-11032, 2020.
[31] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances neural information processing systems, 32, 2019.
[32] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018
12

---

[33] Grant Rotskoff and Eric Hunden-Eijnded Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75(9):1889-1935 2022.
[34] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):725-752, 2020.
[35] Dyego Araūjo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.00193 2019.
[36] Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint 1rXiv:192..02860 02880, 2019.
[37] Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXlv:2001.11443 2020.
[38] Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean-field framework for over-parameterized deep neural networks. In Conference on learning theory, pages 1887-1936. PMLR, 2021.
[39] Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctua- tions in mean field neural networks, arXiv preprint arXiv:2304.03448, 2023.
[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[41] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016. British Machine Vision Association, 2016.
[42] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1492-1500, 2017.
[43] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint 1r0X.vv:900.000
2019.
[44] Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research, 13(1):795-88) 8288, 2012.
[45] Simon Korblithh, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International Conference on Machine Learning, pages 3519-3522 PMLR, 2019.
[46] Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some enns. Nature Communications, 14(1):908, 2023.
[47] Adam X Yang, Maxime Robeyns, Edward Milsom, Nandi Schoots, and Laurence Aitchison. A theory of representation learning in deep neural networks gives deep generalisation of kernel methods. arXiv preprint arXiv:2108 (0088/3300 10097, 2021.
[48] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[49] Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. 8th International Con- ference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL httl:tttts:://ppeeeeeii wttpee/wwwww.eeeeeouum///daaaeee met.:/forrumm?ddsggisstttr
[50] Philip M Long. Properties of the after kernel. arXiv preprint
2021.
13

---

[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 http://pwwreveeeeee.eefrrrrnrnra
[52]
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry, Trak: Attributing model behavior at scale. arXiv preprint arxiv:2303. 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, 13:22217-222335 2020.
14

---

A Experimental Details
A.1 MIPs
In Figure used 3-layer MLP learning Gegenbauer polynomial (2(3.z) in D = 5 dimensions.
Here B was randomly chosen unit vector in IR. D. We implemented μP parameterization by hand.
The output layer of the network was rescaled by au / N, consistent with p.P. We chose og = 1000
to put us in the lazy regime. We set the the learning rate to be 511(( + ag)) General arguments based on kernel scale indicate that the learning rate should be scaled as 0gg" at large a.
In Figure 14 used 3-layer MLP learning Gegenbauer polynomial Q2 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 (RandomHorizontalllill)) in pytorch). As base network (for μP) used ResNet18 where BatchNorm
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 (aa we used the above setup, but with learning rate of 0.01 and much higher batch size of 2000, so to replicate the edge of stability phenomenon [7] which only occurs at high batch sizes. For Figure 6bbl we used 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 ResNets trained on CIFAR- 10 for several epochs. The architecture is 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- k dataset consisting of 220 1048576 labeled images and 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, AutoAugmentPolicyy IMAGENET.
We again used the ResNet- 18 architecture with μP parameterization relative to the ResNet- 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 bb) and (b)) were trained using the Adam optimizer with the following learning rate schedule: linear warm-up for 0.5 epochs from learning rate 8 x 10-5 8 x 10-3, followed by cosine decay over 49.5 epochs to 8 x 10-5
A.3 Language
A.3.1 Wikitext-103 Language Modeling
For all Wikitext-103 tasks, we adopted the pP transformer as defined in the jPP package [33. In the plots shown in the main text, we used a depth-4 transformer, with dmodei- de, d, = N and dsrn = 4N. We performed single pass through the train set in order to stay in the realistic online regime. We used masked language modeling with sequence length S at varying input sequence lengths S. For Figure d) we used the x attention matrix of an S 128 transformer. In Figure 5) we used the attention matrix of an S = 35 transformer. We chose this different length simply to
15

---

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, dmodde).
We used the Adam optimizer with a learning rate of 0.0001. We also ran the same configuration with SGD and learning rate of 0.5 and observed the same behavior. See section B for further plots and details.
For figure used the Wikitext-103 validation set in order to measure the evolution of the predictions on masked logits. ff), D), we averaged the mean squared error from the widest transformer by using 100 test points.
A.3.2 C4 Language Modelling
Figure (b) trained with base network being 125m parameter transformer model on
2.5 billion tokens using the Mosaic ML's LLM codebase (https://web. http::://eeb.arccive..oooo www.iaa..c...eeee..e...eo org/w//b
See http:://ww...ccooooooo/////////2
com poosaicm// wwwmppels//m//yaa///////ooo///ee http://aaaaaaoaaaaoooooo//////// yam for the full hyperpa-
rameter details. We were limited by time and computational details of the C4 transformer model.
resouress resources in our ability to explore further
B Further Plots of Convergence
10
Inkial
Final (20 spachs)
Test Error
Error
258
512
ptn)
10-2
N
33
1
64
N 120
N 254
512
10-4
Epochs
(a) Error Consistency
(b) Last Layer Preactivation Consistency
Inttal
1- C-CAA 112)
10:2
Final (20 epoch)
Init Final
10a
2x100
N
(c) Last Layer Kernel Consistency
(d) Kernel Alignment vs
Figure 10: Training CIFAR-10 10 in 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 the CNN for a randomly chosen data point are Gaussian (as expected) and very consistent across model widths N. To obtain histograms train an ensemble of = independently initialized networks concatenate activation patterns across members of the ensemble. After 20 epochs of training (models around 95% accuracy), the preactivation distributions for the same data point have become non-Gaussian (consistent with infinite width theory) but still remarkably consistent for large widths. (c) The final layer's feature kemel at initialization shows very little structure, but (d) after training networks all widths converge to similar kemelss The plot in (d) compares ensemble averaged kernels with the N 512 ensembled kernel.
16

---

2=2
==3
1=4
N
N 128
256
= 512
Figure 11: Convergence of layerwise representations in each layer (block) of the ResNet-18 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 ResNet learning CIFAR- 10 and going over multiple passes of the dataset. In Figure I0 we plot 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 exhibit pathologies after sufficiently many epochs, related to the discussion in section 3.2
Next in figure across layers for the CIFAR-5M ResNet after a pass through the data. The larger resolution allows see that the fine-grained details in the structure of the Gram matrix are consistent across widths.
show nigher-resolution plot of the kernel Gram matrices across widths and
B.2 Language
Next, in Figure 12 we create an analog the language column of Figure 3 this time for jP 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.
17

---

01.0
M
66
Erga neadRt
c-Starg
0.5
E01 1
Tiet
5-522
C-4055
0.
15P
33P
00
36
(a) Training Loss
(b) Correct logit value
(c) RMSE from widest
Figure 12: An analog of Figure B for pP transformers trained with SGD. a) Training loss. It is interesting that in 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 and SP (Standard Parameterization)
There are several detailed discussions about μP vs SP scaling [2] 15 10 [4]. The aim of this section is to simply give accessible and conceptual overview of their distinction, as well motivation for μP from the perspective of keeping features moving in time even at infinite width.
There several equivalent ways of parameterizing neural networks that give rise to the same dy- namical effects, whether in u-parameterization or standard parameterization. We give the definitions in the case of single-output werd-forward network and demonstrate that SP and μP give rise to ((N-1/2), (1) feature movement at initialization, respectively.
Generalizations to other architectures (ResNets, Transformers) are straightforward. For detailed discussion see and also 14
C.1 SP
We assume all hidden layers have equal width N. Let the input space have dimension D. Let be the index of the training point in the dataset. At each layer É, the pre-activation h in layer + is given by
++1 W o(h)),
(2)
where ó is an element-wise non-lincarity, often taken to be the ReLU function. Here the 1/2 out front allows h to be ((1) initialization N by the law of large numbers. The output of the network fg is then given by:
f
a VN
w (hh,)))
Here again the N-1/2 scaling again yields that fr will be (11) as N + but we will keep explicit as plays an important role in distinguishing the parameterizations. is the laziness parameter identified in 31). The change in the function is given by
co. In SP, a: is taken to be 1,
df d-=-- Σ Kyvl' (fu.yu). dt
---1
Here K₁n T Valu is the NTK gram matrix. y" the true label. l is the loss function (e.g. MSE or crossentropy) and f' is its derivative with respect to the first argument. The NTK is easily seen be order a² and ! is order 1 at small a. In order to have the change in the function be (1)
we set 17 ?
Using the chain rule, directly that the pre-activations evolve as 10
32
dh
Yo
(5)
dt
av N
18

---

Thus, at large N and a 1 the pre- pre-activations of this network evolve as ((N-1/2). Consequently, infinite width the feature do not evolve and infinitely wide networks in standard parameterization become kernel machines with the static and iniiializtioonindepeee
infinite- width NTK.
In many machine learning libraries, the factors of 1/vN are not explicitly placed front of each multiplication with the weight matrices. Rather, the weight matrices themselves are drawn from distribution W ~ (0, N 1). Although this gives identical forward pass, this changes the N scaling of the gradients in the backward pass. As long the learning rate is appropriately rescaled account for this, the dynamics are equivalent to the SP parameterization discussed above.
C.2 aP
One of the simplest ways define the u-parameterization is to take 1/VN. This implies that we simply replace the final layer of the network by:
f ---" w (h))))
(6)
As the prior analysis shows, in order to have dff./dtt dt be (1) at initialization, we take rj = Nooo, so the learning rate in this definition scales extensively with N. In this setting, we now have that
dh di
(().
In 12 [5]
gives an equivalent definition μPP that gives rise to the same dynamics but keeps the learning rate to be (1)) We this version of μP in the experiments that run, simply because that is what is used in the package 13. Consequently, our learning rate does not need to changed width grows.
D Importance of Parameterization
2.0
128
Kernel Alignment Increase
0.0020
128 param.
256
256
N 512
1.8-
0.00115
1024
1024
0.0010
Test
0.0005
Epochs
Epochs
Figure 13: (a-b) An experiment CIFAR-10 with depth 4 CNN (no layernorm) showing that the finite width consistency at the range of widths N considered this work is special for pP and is not observed in NTK parameterization. Width-consistent kernel regime behavior requires 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 μPP networks is that feature learning remains (1))
In this section, we aim to illustrate that many of the claims width consistency at practical widths in fact contingent using the mean feld//-parameeterizaioo.. In Figure [3 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 way so that their N = 128 dynamics perfectly coincide. We plot both test eтror and the kernel alignment between the NTK and the target function. While both parameterizations will eventually converge 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 because the rate feature learning not held constant across widths and decreases width grows. Wider networks in uP train slightly faster.
19

---

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 dynamies (39]
E
Further Studies of μP versus SP
E.1 MLPs
In figure 14] we show 3-layer MLP learning quadratic polynomial. In subfigure a) use batch size of 10 and learning rate going as 1 = 5N/(1 + ag) with dg = 1. The output layer is scaled as an/ /N,,,, putting us in the rich regime. The learning rate has been picked to be nearly as large as possible this batch size in order to maximize the large loss curve fluctuations yielded by large learning rate effects. In subfigure <) we do not rescale the output layer, and have a width iidth-indeppeet learning going as 1 50/(1 + az) with og 1. This puts is in the large-learning rate regime for a standard parameterized network. See Appendix [ 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 tfgge-o--stbbllle size analog. We plot the absolute difference from the widest network in subfigure b) to highlight the strong agreement across widths.
effect reported in [7], the large oscillations may be similar to small batch
In subfigure c), we have the same network but standard parameterization. The narrower networks now learn features more quickly, leading to inconsistent dynamics across widths.
Mostte stoeeece df tow frem 4800
MI
22
22-
Tan
kc-
20-9
Ni-258
Mocces
R1-2008
11-2232
Strpe of 500
Stople of 503
Figure 14: 3-Layer MLP learning a quadratic polynomial y Q2(3 x) 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 For jPP networks, the loss curves match across widths, even accounting for fluctuations due to batch size large learning rate the widest network. c) The same network but in standard parameterization. Dynamics are no longer consistent across widths, and wider networks approach lazy limit.
consistent across large widths. a)
effects. b) Plot of the difference in training error from
E.2
Vision
Next, we focus on a vision task and compare the large learning rate small batch size effects in SP to the μP parameterized network in Figure 6a| By contrast that figure, we see significantly different dynamics and batch variation across widths. In Figure 15 a) we plot the early time behavior of 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 [5] the sharp spikes in the loss function due large learning rate effects become substantially different across widths. Indeed, in SP some widths converge for given learning rate while others not. This trend has already been well-studied in [55.. We again stress that our observation is that not only are the final losses similar across widths in pP (as observed in (15)), but that the individual batch and large learning rate fluctuations agree across widths early times jP as well.
20

---

60
200
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 explosions that happen different times and scales across widths.
E,3 Language
Finally, we present complementary set of figures to those in the right columns of Figures 3 and for transformers of the same architecture on Wikitext- 103 but standard parameterization.
g60
Anuna Ao
Erga nEvERt
Tua 5o01 1A
Hidd
E-128 1rg
N--32
Rr-2224
Wi-4006
0.0
3.2
044
0,5
0.5
1.3
0.0
3.2
0.5
0.0
10
(a) Training Loss
(b) Correct logit value
(c) RMSE from widest
An analog of Figure 3 SP transformers trained with Adam. a) Training loss. large widths, the learning rate chosen is too big for the network to properly learn, and the loss is flat. This consistent with what observed in [15] the optimal learning rate in SP changes with width. b) Value placed the correct logit for specific masked token. c) RMSE of correct logit value from the widest network. both these plots, the monotonic behavior across width evident under the μP parameterization is violated. Even after discarding the networks that do not converge under SP, the behavior remains non-monotonic across width.
Figure 16:
Bias- Variance Decompositions over Initializations
In this section we explicitly define what mean by initilization bias and initialization variance. Following (30). we consider the trained neural network function fo- (ær) to depend on the training set D (including inputs ær outputs y, and possible label noise well the initial parameters 0g- Here, fn are the initial parameters and 6* 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 (z 3/) By contrast, our paper focuses on the variance due to the initial parameters Og from which training begins. In the overparameterized regime with more parameters than data-points and relatively little label noise, this has been shown be the dominant
source of variance for neural networks 22 29 28)
Consequently, our definition of the bias of neural network is given by
f(x) Ee, (0)-
(8)
The bias can be approximated by averaging sufficiently large ensemble of neural networks over initialization seeds. Fach network the ensemble trained on the same dataset in the same batch order using the same optimizer.
21

---

The variance of the neural network predictor is given by:
Vare, = Ee [[fo (0)) ( f(x))
(9)
The variance of ensembles of given network with independent initialization seeds 0g given by
)
2
Edp [ -
Vars. Σ (z))
f(x))
(10)
As E - oo, 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 1/E. In practice, see that comparable decay rates 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 I epoch for Wikitext 103). We fit power laws of the form Loss CN + D where C, 3, D are fit to the data using the "scipy.optimize" function. The resulting fits are provided in Figure 17 We find an excellent power law fit, with R above .99.
The existence of such power law behavior across widths provides further evidence of the networks ap- proaching a well-defined inifinite width limit. One could imagine that the differences between succes- sive networks might get smaller (as Figure 4) but that there is well-defined limit as N + oo, 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 pro- vides empirical evidence that we expect convergene as N + 0o.
Loss CN-P+ D
Data, Wiittxt N A1-042
10a
Data, Imagenet
11.2 x
66-10-
4
3x10
x
Additionally, given that the observed power laws N-9 are task-dependent and significantly different from N-1 suggests suggesss that models at late time are not well described by perturbation theory around the infi- nite width mean field limit, which predicts universal exponent of 3 = 1 [54] 39]. This motivates novel theoretical descriptions of finite width mean field learning networks at late time which can capture task-dependent exponents.
10a
10a
Figure 17: Train losses of Wikitext and Ima- genet models are well-fit by power laws.
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 MI.P in a parameterization equivalent to yP
1 N ((h))
1
1
fp
th
f+1 h
W (h h/,
W Z..-
(11)
VD
We will consider these networks trained from random Gaussian initialization of the weights so that O Vec(w= ..., Wo} follows 0 ~ (O I) at initialization. This network is then trained with some gradient based optimizer, leading to dynamical predictions ft(() and dynamical preactivations h (1). Because of the random initialization of weights, the outputs of the network and the precise preactivations are random variables. However, at infinite width N - oc, a dramatic simplification of the dynamics occurs.
22

---

H.1 The Infinite Width/Mean Field Limit
The predictions (1) and internal representations of infinite width limit of neural networks admit description in terms of non-random intializationniideeeeneenn dynamical feature kernels &e f (t, 8) and gradient kernels Ge (1,8) defined as
1 & (1, N (()(()) ((h)()) G (l, (1, 5)
N g/(1)) (s),
(12)
where μr, index data points and f, s index training time and gj. gf. (t) = N Uh gradient signals [12 14]. Further, all preactivation vectors h hf,(t) E IR have entries that become iid draws from (potentially non-Gaussian) single site density p(h), which converges
oh are back propaapodd
N Σ
8(h h4) - p(h),
(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 of Ti (t, s) = p(h,, (t), & (s)) (t)) (t))o(( ()) (t) (s).
H.2 Finite Width Effects
At finite width, the internal kernels {Φf (t, G (t, 8)} and predictions (t) of the model become initialization and width aedth-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 (1/NN) for width N (or 1/dmodel for transformer) 39]. Further, the ensemble averaged values for the predictions (fx(t)) and kernels (g (t, s)) differ asymptotically from their infinite width values by ((N-1). Both of these two leading order effeets 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 asymptotic series in powers of N-1 [39].
H.2.1 Trainability at Finite Size
The O(N-1) correction feature and gradient kernels can lead non-trivial corrections to the loss dynamics. Working in continuous time, we can define the neural tangent kernel (NTK) as Ky (t) >, Gtt ff (t, where base cases are & oh (t, s) 1 -h Vy and G s) 1.
Following the approximation to online dynamics with MSE loss in Section 4 flow the average dynamical NTK
we consider gradient
d ( - (K(I)) (!) = A(I) = T exp (-f [-]'ds (K())) y. di
(14)
where T is the time-ordering operator. We now consider the leading correction to the average NTK around infinite width (K(1) = K (1) + x (²(()) + O(N-2). With this correction, we see that the dynamics of errors A
A (t) = T exp (- (-f" ds 1 J. [s) y. N
(15)
The fact that the N K correction integrated over time and placed in the matrix exponential indicates that small corrections 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(L) defined
d T(t) (K(t) T(t), T(0) I, R(t) log T()). dt
(16)
The solution to this can be used to construct the errors at later time A (t) exp (R(t)) y.
23

---

I The empirical sufficiency of ensembling only a few times
Our experiments rely on ensembles of small numbers of neural networks analyze the bias compo- nent 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 I8 illustrates this in the setting of ResNets trained on ImageNet. In Figure I8ja, a), we show that the reduction in variance over initializations due ensembling rapidly plateaus as soon as the ensemble size reaches E = In Figure [8))), 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 E (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 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 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 E {128, 256} compared N 512.
22
322
322
2.8
22
0.3
Ensemble size
Base withh 00
(a) Test loss vs. # ensembles
(b) Test loss vs. width
(c) Test loss for subtracted networks
Figure 18: ResNet- PosNeet-1/IIiae 1/ImageeNe loss as a function a) ensemble size or b) width after epochs of training with heavy data augmentation. c) Comparing original network with network whose output has been set to zero at the start of training.
J
Offline Training
Figure I9 depiics the loss curve for ConyNeXi-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 widte networks. Width consistency subsequently erodes.
8.0
4.0
Toan Lot
1.0
6.5
1000
10000
Train Steps
Figure 19: Consistency of loss curves across widths in the beginning and separation loss becomes sufficiently small in offline learning.
24

---

K Use of Compute
For most experiments, we used Nvidia A100 SXM4 40GB and 80 GB GPUs on academic cluster.
For the Wikitext- 103 tasks, each width included ensembles loaded onto an A100 GPU that ran for range between Accounting for sweeps over different sequence lengths, optimizers, and parameterizations, this corresponds to about 50 A100-days.
days. For each sweep over widths this corresponds to about 8 A A100-days.
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 10 ResNet experiments in Figure T0 were done using a total of less than A A100-day of compute across all widths and ensembles.
The ImageNet ResNet experiments vectorize training over between to four -iine neural networks on one A100 GPU. Each experiment training collection of networks for 30 epochs takes between one three A A100-days. Overall, these experiments expended roughly 30 A A100-days.
For the CIFAR-5m experiments in Figure and 3 across widths, it required a few hours of A100 GPU. For Figure 7 8 these were ensembled across multiple runs, these required close to 1-2 days of A100-GPUs. Figure 6a was just run for few 100 steps of the training, so didn't use much compute power.
25


Reviews:
[{"review_id": "4G8ulRHLDC", "reviewer": "Reviewer_ddkm", "summary": "In this work, the authors conducted an empirical investigation into the consistency across widths of the training dynamics of neural network models in the maximum-update scaling. Specifically, they demonstrated through experiments on both vision and language tasks that many behaviors in the online learning dynamics are consistent at large widths, including the loss values, single predictions, feature kernels and large-learning-rate phenomena, and that ensembling improves the performance. Deviations from the large-width behavior are reported in the offline learning setting. Finally, the authors argued through heuristics and experiments that the finite-width bias is largely due to the deviation in the tangent kernel's eigenvectors.", "questions": "(Please see the comments above.)", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Given the number of theoretical results concerning over-parameterized neural networks and their infinite-width limits, the question of the finite-width deviation from the infinite-width limit is an important one, especially under the feature-learning regimes. To my knowledge, this is the first work to put forth an empirical investigation of the large-width consistency of neural networks in the maximum-update parameterization. The authors of the paper performed a wide range of numerical experiments on various tasks with different neural network models that cover several relevant aspects of the model training behavior.", "weaknesses": "Note that for feature-learning neural networks in an alternative (mean-field) parameterization, both empirical and theoretical studies of similar purposes already exist, with examples including Nguyen (2019), Chen et al. (2020) and Geiger et al. (2020). More comments below regarding parameterizations.\n\nThe authors argued based on Figures 2 and 7 that, as the widths of the model increase, the training performance of the single network converges in the online learning setting but not in the offline setting, which I don't find fully convincing. The y-axis uses a linear scale in Figure 2 and logarithmic scale in Figure 7, and hence the discrepancy of the losses across different widths in Figure 7 (which appears to be quite large late in training under the logarithmic scale) may not actually be much more significant than in Figure 2. In addition, it is also hard to rule out the possibility that the width being used is not yet large enough to approximate the large-width limit (after all, 512 is still small compared to the size of the datasets).\n\nThe original proposal of the maximum-update parameterization does not include normalization layers, and given their role in preserving the scale of the hidden-layer features, one would expect their effect to interact with the initialization scale of the parameters. The authors only mentioned in Appendix A.2.1 that LayerNorm was used in the experiments on CIFAR-5m and did not comment on whether normalization layers are used in the other experiments. I think the authors should discuss in detail whether normalization layers were included in the other experiments and, if so, why the phenomena observed in this work can be truly associated with the maximum-update parameterization.\n\nReferences:\n\n1. Nguyen, \"Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks\", arXiv 2019\n\n2. Chen et al., \"A Dynamical Central Limit Theorem for Shallow Neural Networks\", NeurIPS 2020\n\n3. Geiger et al., \"Scaling description of generalization with number of parameters in deep learning.\" Journal of Statistical Mechanics: Theory and Experiment, 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for suggesting these important related works. We will be sure to include a proper comparison of our study to these previous papers. Below, we provide a brief comparison of our results to each prior work.\n\n#### Comparison to prior works\n1. First [1] Chen et al study two layer neural networks in the mean field parameterization and establish a central limit theorem where the mean of the neural network predictor dynamics is the infinite width limit and the fluctuations scale inversely with width. However, this finding does not carry over to deeper networks, where the average of finite width networks deviates from the infinite width limit (what we call the bias correction of finite width). This gap between ensembled narrow networks and wider networks observed in our work demonstrates that ensembling many finite width networks is not equivalent to the performance of an infinite-width deep network, even with infinite ensembles. Chen et al also provide experiments showing lawful fluctuations that scale as $1/N$ but have similar dynamics in student-teacher settings and a simple non-planted two dimensional target function. We will be sure to cite this paper as evidence that fluctuations and average dynamics can be consistent for large but finite $N$. \n2. Second [2] Nguyen et al study a slightly different parameterization which also gives an infinite width feature learning limit. To see the difference between our parameterization and theirs, note the factor of $1/N$ in equation 7 of [2] for the intermediate layer, which would be $1/\\sqrt{N}$ in our paper. Nguyen et al do provide experiments which show width consistent dynamics in their parameterization for MNIST and CIFAR-10 in fully connected networks with a handful of layers $\\sim 5$. This parameterization behaves differently at infinite width compared to the parameterization we study. Concretely, in another work (Nguyen & Pham 2020 arXiv:2001.11443), Nguyen and Pham show that iid random initialization leads to a degenerate limit for any depth $L \\geq 4$, which makes it impossible for successful finite width training from iid init to be consistent with infinite width networks. On the other hand, the $\\mu$P infinite width limit is nondegenerate for any finite depth $L$ with iid initialization. This makes it *possible* that finite width networks could be consistent with their infinite width dynamics and motivates an empirical study of finite width consistency. Further, our work provides experiments for deeper networks with a variety of architectures and datasets including Residual CNNs (depth $\\sim 20$) and transformers ($\\sim 4$ encoder blocks which is like $\\sim 12$ total layers) and we document consistency for other relevant deep learning phenomena (feature kernels, preactivation distributions, large learning rate effects etc). \n3. Geiger et al [3] provide a detailed study of NNs trained on NTK parameterization, where the large width limit is given by a kernel method. Their study also documents interesting fluctuation effects from initialization variance and have experiments on MNIST and CIFAR-10, convincingly arguing that corrections to infinite width scale as $1/N$ due to finite size fluctuations. Our study is different from this study in that we provide $\\mu$P networks which have a feature learning infinite width limit and we focus on dynamics of many observables during learning, rather than the error curve for fully trained networks. A different work from Geiger and colleagues (which we currently cite) does study the dynamics as a function of feature learning strength and characterize a parameterization that is equivalent to $\\mu$P as a limit where width consistent feature learning is obtained (Geiger et al 2019 arXiv:1906.08034).\n\n#### Layernorm\nWe thank the reviewer for this question and allowing us a chance to clarify our use of LayerNorm. In the original submission, the ResNets on CIFAR-5m and ImageNet and the transformers trained on Wiki-text all utilized LayerNorm. The experiments on MLPs in Section 4 do not have LayerNorm.\n\nLayernorm does not prevent a well defined infinite width limit in $\\mu$P. Several works have established the infinite width signal propagation of networks with Layernorm (Doshi et al 2021 arXiv:2111.12143, Yang 2020 arXiv:2006.14548, Yang & Hu 2021 arXiv:2011.14522). The key idea is that preactivation vectors (before layer norm) $h^\\ell \\in \\mathbb{R}^N$ will have means $u^\\ell = \\frac{1}{N} \\sum_{i=1}^N h^\\ell_{i}$ and variances $V^{\\ell} = \\frac{1}{N} \\sum_{i=1}^N ( h^\\ell_{i} )^2 - ( u^\\ell )^2$ that concentrate in the $N \\to\\infty$ limit. These objects $\\bar{h}, v$ are covered by Yang's Master Theorem (Theorem 7.4 here arXiv:2011.14522).  Thus at infinite width, LayerNorm will merely shift and divide each preactivation vector by a deterministic value $h^\\ell_{\\mu i}  \\to \\left(h^\\ell_{\\mu i} - \\bar{h}^\\ell_{\\mu} \\right) / \\sqrt{v^\\ell_\\mu}$. Backpropagation signals can similarly be reasoned about. Frameworks like Tensor Programs (Yang & Hu 2021 arXiv:2011.14522) or Dynamical Mean Field Theory (Bordelon & Pehlevan 2022 arXiv:2205.09653) could be used to more carefully demonstrate how $u^\\ell_\\mu(t), v^\\ell_\\mu(t)$ will concentrate and evolve over training time $t$. We will comment on this in the Appendix.\n\nIn addition, to further test that the width consistency we find is not an artifact of LayerNorm, but is also visible in $\\mu$P networks without layernorm, we provide an example experiment on CIFAR-5m. In the rebuttal PDF we show an experiment of many networks trained without LayerNorm showing width consistency (Figure 4)."}, {"review_id": "nLcXlScEdZ", "reviewer": "Reviewer_gHF9", "summary": "This is a primarily empirical paper studying the agreement between real, finite-width neural networks and their limiting behavior under infinite width. Rather than using the neural tangent kernel as the descriptor of infinite-width behavior, the paper focuses on a different parameterization called feature-learning networks, which vary from standard-parameterization networks in that intermediate layer weights do not remain fixed even in the infinite-width limit. For these feature-learning networks trained in the online setting (without repeating training data) across a mixture of architectures and vision and language tasks, the paper demonstrates that large but finite-width networks tend to agree across a wide variety of metrics including overall loss, per-sample predictions, and the distribution of learned weights. However, if networks are too narrow, or are trained in the offline setting (when data is repeated during training) this agreement can break down. This gap between finite and infinite-width performance is then studied using the empirical NTK (at initialization) as well as the empirical after-kernel, and the results suggest that finite width induces deviation in the kernel eigenvectors (but not its eigenvalues) compared to the infinite-width limit.", "questions": "- Section 3 describes ensembling of finite-width models as a way to reduce variance, but shows that ensembled smaller models are still worse than a single larger model. How did you decide how many smaller models to ensemble? In particular, I wonder if there is some combinatorial flavor here (inspired by https://arxiv.org/abs/1803.03635) where an exponential number of smaller models is needed.\n- In Section 3.2, what motivates the use of these two datasets? In particular, I would think comparisons of the effect of label noise would be easier to see if using the same dataset with vs without label noise, rather than two different datasets where one of them happens to be noisy.\n- In Section 4, why use NTK and after-kernel? I presume this is because there are not similarly computable limits for feature-learning networks? In general I wonder why the paper uses large but finite models as a proxy for infinite-width networks, if feature-learning limits are well understood. An explanation of this reasoning and more detail on related work for what is known about the infinite-width limits of feature learning networks would be valuable.", "limitations": "Limitations of the experiments are not directly discussed, and the paper would benefit from such a discussion, perhaps in the conclusion. One limitation would be the lack of theoretical or direct modeling of the infinite-width limit, and the use of large-width networks as a proxy. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I really appreciate the thorough and systematic experimentation. Deep learning theory benefits from the ability to model infinite-width networks, but it’s important to know how well and under what circumstances such infinite-width limits translate into finite-width practical behaviors. In particular, I appreciate that this paper includes a mixture of convolutional and attention-based neural networks on a mixture of vision and language tasks, and studies finite and infinite-width agreement across loss, per-sample predictions, and network weights. It also finds some interesting and useful boundaries for when we can expect agreement to hold versus when finite-width networks behave differently, and takes steps to explain why.", "weaknesses": "Most of the weaknesses I’ll describe should be viewed as suggestions to improve the exposition of the paper, and in some cases suggestions to make the experiments clearer or more compelling. I tried to separate suggestions/complaints into this section and clarification questions into the “questions” section, but there isn’t always a perfect distinction.\n- The abstract uses the term “feature-learning neural networks” without definition. Thought it’s defined early on in the introduction, many potential readers may be confused by the abstract before they see the definition in the introduction.\n- There are a decent number of typos and small grammatical mistakes, though they don’t impede comprehehension. For example, Figure 1 caption says “Loss curves…are nearly to identical at large widths.” Line 37 says “We put all of our in mu-parameterization”—similar missing-word type typos appear throughout the text.\n- Essentially all of the figures require the reader to zoom in to understand what is being shown. Figures should be fully interpretable without zooming in, e.g. for readers who use a printed copy. Axis labels, legends, tick labels, and line widths all need to be substantially larger. \n- Lines 122-125 define a bunch of notation that is (1) confusing, especially for readers unfamiliar with Transformers, and (2) is never referenced again. Instead, figure 2 uses “d” without any subscript, which is not defined. It would be preferable to define (only) the necessary notation, and push other details to the supplement.\n- Figure 3 is a bit confusing. It’s not clear if the first row is literally showing results from a single test point, or if it’s an average over test points. I’m also not sure why these figures only consider the correct logit value as a metric of agreement; I would have thought a more natural comparison to be some metric of model agreement across predictions, such as overlap in where mistakes are made (used in https://arxiv.org/pdf/1905.12580.pdf).\n- Figures 4c and 4f and 8c and 8f all plot 1-{some quantity} on the y axis, whereas in the textual description trends are described with respect to {some quantity}. This is needlessly confusing.\n- Line 157 describes figure 5a as growing “steadily to a final value that it then fluctuates around” but in the figure I don’t really see a stable final value with fluctuations; rather it seems the trend continues to increase even after fluctuations arise. This is a minor point, but was a bit confusing since it otherwise seems the main purpose of the plot is to visualize when the fluctuations begin.\n- Figure 7: Please label the x axis with epochs either instead of or ideally in addition to training steps. The purpose of the figure is largely to investigate the effects of reusing training data from one epoch to the next, so it’s important to see when data is first reused. For example, in subfigures (a) and (b) I suspect this happens around step 10000 when the curves diverge, but if this is the case I wonder why not include more training steps so that we can see the deviations more clearly?\n- Equation 1 seems to be describing the NTK. I wonder why the preceding paragraph that describes its derivation doesn’t mention NTK? As I read that paragraph I kept thinking “this sounds just like NTK” but I wasn’t sure why the definitions were present unless either it’s different in some way from NTK (in which case, please specify the difference directly) or it’s a refresher for readers who aren’t familiar with NTK (in which case, please say so to avoid ambiguity).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful comments and suggestions. Below we address the various weaknesses and questions which arose in the review.\n\n### Weaknesses\n\n1. We will include a sentence in the abstract specifying that we mean networks in $\\mu$P. We agree that we should define what we mean by \"consistency of feature learning networks\" earlier on.\n2. Thank you for finding these incomplete sentences and grammatical errors. We will comb through the text and check for these errors. \n3. We will increase fontsize and figure size using the additional space, provided that our paper is accepted.\n4. We will define $d$ as the embedding dimension used for the queries and key vectors. We wanted to be explicit what our transformer sizes were for readers whose primary interest was language models. We can move this discussion to the appendix if it confuses or detracts. The $d$ in Fig 2c is indeed confusing and we will replace it with $N$.\n5. Figure 3 is a *single test example* from the test set but the curve is averaged over the random initial networks. We will clarify this in the new version.\n6. We will try avoiding plotting (1 - {quantity}) in our plots. Either by redefining the metric so it is just {quantity} on the y axis or just plotting the current {quantity}. We preferred 1-CKA to CKA since it lets you see the scaling law of convergence of the kernels on a log-log plot.\n7. In plots involving multiple passes over the data, we will make the $x$ axis epochs. We agree this will allow us to see the effect of multiple passes over the same data. In the plots in the paper and the attached pdf data actually starts to repeat much earlier, at 400 steps! (though there is data augmentation) So 10s of epochs are needed for overfitting to emerge. \n8. Equation 1 is definitely the NTK! We are sorry this was not stated clearly, we only say \"the kernel\" but what we really mean is the \"neural tangent kernel\" for the NN. \n\n### Questions\n1. We experimented with different ensemble numbers and verified that the results were stable under addition of new ensembles. See Figure 5 left panel in attached rebuttal PDF for Imagenet experiment with varying number of ensembles. The curves quickly converge in $E$. For MSE loss it is easy to theoretically show that the loss decays as $B + S/E$ where $B$ and $S$ are fixed constants.\n2. Both of the experiments in Section 3.2 are just subsamples of CIFAR-5M. They were both chosen so that the data could be passed over multiple times. For the noisy data, we saw a more dramatic overfitting effect if the dataset was smaller. We can provide plots for other subsamples if the reviewer is interested. \n3. In section 4, we look at the final NTK measured from data since we do not have a computationally tractable way to compute the infinite width limit or its finite size deviations. Empirically other works have shown that the after kernel does seem to capture important properties of the final model (Long 2021 arXiv:2105.10585). We will add additional information about what is known of the infinite width feature learning limit and why it is challenging to compute."}, {"review_id": "4oOhJBBkHA", "reviewer": "Reviewer_g9eH", "summary": "The paper investigates the impact of width on the dynamics of feature-learning neural networks with both vision and language tasks. The empirical results demonstrate that during the early stage of training or during the entire training process for easy tasks, wide neural networks trained on online data not only exhibit identical loss curves but also display consistent point-wise test predictions. Nevertheless, for more challenging tasks, in the later stage of training, deviations across finite widths become more pronounced. The authors then decompose the deviation to bias and variance and show that the variance scales inversely with width. Furthermore, the authors demonstrate that ensembles of narrow networks can have inferior performances compared with a single wide network.", "questions": "In section 3, the authors decompose the deviation into bias and variance, how did the authors quantify the two terms? It seems like bias is represented by the training loss, and variance is quantified by the improvement brought by ensemble?\n", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized and clearly presented.\n2. All the experimental results are clearly motivated and effectively support the claims made in the paper. \n3. The supplementary material is comprehensive, providing extensive details and codes for numerous experiments conducted in the study, enhancing reproducibility and transparency.", "weaknesses": "he reviewer's major concern about the paper is the potential implications of the paper. While the empirical results presented in the study are well-documented, it remains challenging to discern how this work motivates or contributes to the broader research field. In the conclusion section, the authors assert that the findings \"motivate the applicability of infinite-width feature-learning models in reasoning about large-scale models trained on real-world data.\" However, the empirical results within the paper indicate that deviations caused by varying widths amplify for more complex tasks or longer training time. Moreover, there is no apparent evidence suggesting that increasing width follows the law of diminishing marginal utility. Consequently, the reviewer struggles to grasp how the results presented in the paper effectively support the applicability of infinite-width models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for their comments. We are glad that the paper was received as well-organized with well-motivated results. \nWe would like to address the main concern: the broader applicability of our findings. \n\n* *\"In the conclusion section, the authors assert that the findings 'motivate the applicability of infinite-width feature-learning models in reasoning about large-scale models trained on real-world data.' However, the empirical results within the paper indicate that deviations caused by varying widths amplify for more complex tasks or longer training time.\"*\n\nWe agree that a large part of our results analyze the leading deviations from infinite width behavior rather than showing it holds ubiquitously. We are thus happy to change the final sentence to something that more accurately represents the conclusions of the paper. We have modified the sentence to \n\n\"These studies motivate the study of infinite-width feature-learning models, together with their leading deviations, as a promising avenue for reasoning about the early training periods of large-scale models trained on real-world data.\"\n\nDespite the amplifications at later training times, we have consistently found that at any given fixed training time, wider networks converge to a limit, and sufficiently large networks are consequently quite close to the infinite width limit. The widths required were computationally accessible at late training times even in the realistic settings of ImageNet and Wikitext. Even when the network deviates from the infinite width limit at later times, the leading order deviations can still be systematically studied. Since the writing of this paper, infinite width feature-learning theory was extended to handle finite width deviations [1]. \n\nFurther, for simpler tasks commonly used as benchmarks for deep learning theory, the infinite width limit is quite easily reached for moderate widths and persists as a valid approximation throughout training. \n\n\n\n* *Moreover, there is no apparent evidence suggesting that increasing width follows the law of diminishing marginal utility*\n\nTo address this comment, we have included a set of new plots. Figure 5 in the rebuttal PDF consider the loss of a muP-scaled ResNet-18 model on ImageNet after 2 epochs as a function of $N$. Across all training times we find diminishing returns to increasing the width. We also show diminishing returns as a function of the size of the ensemble in the same panel of figures. \n\nGoing beyond this, we have also added Fig 3 which plots an explicit fit to the loss as a function of width. We fit $C, D, \\beta$ in the equation $L = C N^{-\\beta} + D$ via the scipy curve_fit function. We find an excellent fit to the observed data for both language and vision tasks. This is consistent with diminishing returns of increasing width. We plotted this for a given training step near the end of training, but consistently observe the same power-law trends in $N$ throughout training. \n    \n* *the reviewer struggles to grasp how the results presented in the paper effectively support the applicability of infinite-width models.*\n\nIn Fig 1 of the supplementary material, we plot a network trained on a vision task in both the $\\mu$P and the NTK parameterizations. The NTK trained network does not converge to the infinite width NTK behavior over the given set of widths, while the $\\mu$P networks' loss curves become consistent and closely match. Although infinite-width feature-learning networks are challenging to simulate, it does appear that realistic networks can approach their behavior far more easily than NTK-parameterized networks can reach the NTK. \n\nWe hope that these new studies, together with the refinement of the conclusion, prove to be more convincing.\n\n**Response to Question:**\n\nWe agree that it would be better to be far more explicit in the text about the meaning of the terms \"bias\" and \"variance\". In the final draft, we will always preface them with \"initialization bias\" and \"initialization variance\". \n\nWe will clarify what we mean by *bias* and *variance*. By bias, we merely refer to the deviation between the ensembled predictor and infinite width predictor $(\\left< f_N \\right> - f_\\infty)^2$. For the MSE between $f_N$ and $f_\\infty$, we have the typical bias variance decomposition $\\left< (f_N-f_\\infty)^2 \\right> =  (\\left< f_N \\right> - f_\\infty)^2   + \\left<(f_N - \\left< f_N \\right>)^2 \\right>$.\n    \nWe have added a section in the appendix that now includes the discussion in the above paragraph and adds more detail. We hope this allows for readers to understand how our notions of bias and variance are quantitatively analogous with the standard ones in statistical learning. We also have added a reference to a paper that we believe most clearly elaborates on such a decomposition [4].\n\nIf there are further concrete changes that you would like us to make, or additional experiments to run, we would be happy to include them in the final draft. \n\n[1] Blake Bordelon, Cengiz Pehlevan. Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks. 2023\n\n[2] Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d’Ascoli,  Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization  with number of parameters in deep learning. 2020\n\n[3] Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. 2022\n\n[4] Ben Adlam and Jeffrey Pennington. Understanding double descent requires a fine-grained bias-variance decomposition. 2020\n"}, {"review_id": "0LcpVY38md", "reviewer": "Reviewer_8Tjm", "summary": "The authors study the so-called “feature learning” parameterization of neural networks that is given by the maximal update parameterization (muP). This is in contrast to the NTK parameterization for infinite-widths, which leads to the “lazy learning” regime which is claimed to not learn features, and also in contrast to the standard parameterization for finite-width networks, which the authors also claim (I believe in the Appendix) leads to no feature learning in the large-width limit. The central question of this paper is whether a finite-width (but sufficiently wide) network in the muP setting captures the dynamics of its infinite wide counterpart (which the authors empirically examine by proxy via the largest width network they are able to train, which I’ll refer to as the infinite-proxy). They consider this question in the online setting, where new data is presented in each batch, and the offline setting, where data is recycled over epochs. \n\nThroughout the paper, the authors study the following metrics: training loss, train / test error, logits with respect to the true class for a test sample, attention matrices, top eigenvalues, centered kernel alignment, preactivation neuron distributions, and kernel spectra.\n\nThe central findings of the paper are:\n- Sufficiently wide networks begin to show similar loss, error, and true logit trajectories over the course of training, though this degrades over training on more difficult tasks (e.g. ImageNet vs. CIFAR tasks)\n- You can do a bias-variance decomposition of the finite-width networks by considering ensembles of predictors, and the paper claims that, empirically, sufficiently wide networks don’t suffer much from bias nor variance but narrow networks can suffer from both and that the bias is actually a larger negative effect than variance\n- In offline training, if the networks start to overfit then narrow and wide networks show that bias and variance measured here to not go to 0 (there is a claim that on test the bias and variance goes to 0 but on train it doesn’t, I will defer comments on this to later)\n\nEDIT: Thank you for your thorough response, I continue to be more on the accept side than reject and opt to keep my score.", "questions": "How small of a difference on a metric like cross-entropy loss is needed in order to consider it a negligible difference, or “consistent” as defined in this paper? I imagine this is specific to dataset, model, and other choices. It would help if you gave some idea of this in the paper.\n\nYou could compare the infinite-proxy network (the largest width network that is meant to be a proxy for the infinitely-wide network) to the expected regressor in a simple, synthetic setting (like mixture of gaussians or something like this) in order to see whether the bias of narrow width is like the true bias or it is some new bias-like term you’re proposing, but not exactly equivalent to the actual bias of the model.\n\nOn individual model experiments and true logit experiments, are you averaging errors and losses over multiple different initializations when you plot? The ensemble experiment tells me something about averaging predictions and then computing error / loss, but I think since there is such a high variance with respect to random initialization of the network that the losses, and true logit RMSE and other measurements in most experiments should all be reported as an average over initializations (or test points when applicable), and I just haven’t seen anywhere in the paper or appendix that says that this is what you’re doing, though I may have missed it.\n\nI’m a bit confused about the claim in Section 3.2 that on train the bias/variance don’t go to zero and on test they do so this would indicate benign overfitting. I think it may come from the definition of bias-variance you have given, potentially not aligning to the notion of statistical consistency that would indicate benign overfitting. Could you talk a bit about why you think this is showing benign overfitting, and whether your bias-variance is actually the same as statistical consistency?\n\nIn the first sentence of the second paragraph of Section 4, starting “Concretely, we see that although the eigenvalue spectrum of the …” I find it quite confusing to read, and I can’t help but feel there must be simpler language to get across the point you’re trying to make. Also what is “the task”? Is it fitting the data? This section reads a bit disconnected to me from the rest of the paper and introduces a lot of new terminology and ideas. While I follow the general point, it feels like you’re bringing in a whole second paper’s worth of topics and then trying to compress it all into one page, which makes it feel very rushed and somewhat out of place to me with respect to the prior experiments. I could be wrong, and maybe my fellow reviewers will disagree with my view on this.\n\nI think you’re arguing that the eigenfunctions and how the weight is placed on them is affected by narrow width more significantly than the spectrum itself, and that this would contribute to your notion of bias being large? To me, it feels that you’re throwing a lot of new ideas at the reader in the end of your paper, like “spread out target function power into slower modes” and “putting more of the task into smaller eigenmodes that take longer to be learned” and none of these ideas are motivated or explained anywhere in the paper? Maybe you are assuming that the reader is very familiar with what these phrases would mean but I’ve read a number of papers in this domain and while I can probably make some assumptions about what you’re trying to say I think it’s a lot to ask of the reader to just go with it, unless they are deeply familiar with this specific vernacular.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "There are a lot of experiments and perspectives here to support the claims the authors are making and I don’t think I’ve seen a finite-width study of muP yet, myself. Experiments are also done with MLPs, ResNets, and Transformers, which is an extensive enough set of model families to convince readers that these phenomena are not unique to any one model class. I think the conception of each experiment is nice, and it is clear from the explanation of each plot what it is trying to convince me of. It is clear to me that after a certain width, the finite-width networks begin to act similarly, though there are still minor differences (more on this later). \n\nThe ensemble experiments give rise to a sort of bias-variance decomposition, though I think this is not formally a bias-variance decomposition because the definition of bias here given is with respect to the infinitely-wide version, considered as the largest width network they can train, but I think the infinitely-wide version of the network is not necessarily the optimal predictor (e.g. E[y|x] in a regression setup), unless it is in this setting and I have missed the explanation why? However, the variance aligns to the definition of variance I am familiar with, and in general I still think this is an interesting definition in-and-of-itself, just not perhaps a formal bias-variance decomposition.\n\nI appreciate the thoroughness of everything and overall I find this to be a useful contribution with respect to studying muP networks, something I haven’t really seen before in such a practical setting. I can see through the empirical results that this parameterization does lead to a reasonably stable behavior when the network width is sufficiently large, and it is a useful result for those who want to study and understand the muP dynamics.\n\nOn one hand, I think there is enough here that it will be of use to some people in the community and is worth accepting. On the other hand, I have many comments and questions and I think a second draft and resubmitting to another conference would make this much stronger. I am giving it a weak accept because everything is technically sound and interesting, but I hope the authors consider my comments below.", "weaknesses": "I think it could be misleading to say that you are studying feature-learning networks in the sense that I believe standard parameterization also does lead to “feature learning”, though admittedly I don’t think feature learning itself is a well-defined term. I understand that muP is meant to indicate some non-trivial change in the feature kernel, and in the appendix you argue that standard-parameterization leads to no changes in the feature kernel when width goes -> infinity, doesn’t this still leave room for feature learning in SP in the finite-width (in particular, the narrow-width) regime? I’m not totally sure how you could assess at what point a SP model is doing feature learning or not, but I am open to hearing thoughts on this. Basically, if the point of the paper is to study realism, like realistic scales and realistic models + data then I think it’s too strong to assume that there is no feature learning in realistic SP models. As such, I think it could be more precise to say that you are studying muP models at realistic (finite) scales, and that there are potentially feature-learning dynamics that are not covered in this paper nor encompassed by the muP parameterization. Though I admit I don’t fully know all of the details of muP, so maybe this is just my lack of understanding.\n\nTo be more clear, I think it is obvious once you read the paper that you are only studying muP models, so my prior comment is mostly about the title and abstract, which does not refer to muP and could make a potential reader believe that you are studying, somehow, all feature-learning networks. Maybe the paper could make this point about all feature-learning networks if you compared these muP plots to the same for a modern SP model over a few different widths, perhaps a model that is reported from some recent paper on these datasets (they are very well studied datasets so there will be lots of options to choose from). Is muP outperforming a state-of-the-art SP model? If not, can you claim something about whether there is feature learning or not in that SP model? I understand these questions are not the central focus of your paper, and so they are not affecting my decisions on accept / reject. It is mostly a comment that when the title and abstract says you are studying feature-learning networks this is what I immediately think about and it is something I am wondering as I read your paper, and perhaps having that thought in my head undercuts the value of your paper, to me, while I read it.\n\nI think all of the figures are far too small, especially the fonts. It is hard to read without zooming in on the computer, and on printing it is too small. Especially considering that you want to convince the reader of “consistent” behaviors over increasing width and we are meant to observe this by seeing that loss curves, etc become closer on the higher widths, larger and clearer pictures seem very important to me. I think the small figures undersell the evidence of the claim a bit, it’s not always clear to me that some curves are closer while others are not. In that sense, Figure 3(d-f) gives me the clearest evidence of the claims being true, while many of the rest I really have to stare for a bit and zoom in to see the curves more clearly and convince myself that it’s true.\n\nI also feel that there is so much going on here, it’s not necessarily a bad thing but it’s a bit overwhelming as a paper. The sections just seem like they’re going so quickly through results just stating this figure shows this, now this other figure shows something else, etc etc. I feel that I’m missing intuitions, explanations, contextualization in existing literature and more. If I were deeply involved in muP research then maybe all I need is a clearly stated set of results and a wide variety of experiments to cross-reference, but I think the paper could benefit from more exposition. \n\nPlease see questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed review. We address the concerns and questions below.\n\n1. **Usage of 'bias' does not correspond to usual statistical notion of bias'**: We agree that the way we use bias is not the statistical notion of bias for exactly the reason pointed out by the reviewer: infinitely-wide version of the network is not necessarily the optimal predictor. We will clarify this. We do note that it does correspond to the bias if we take the target function to be the output of the infinitely-wide $\\mu$P network. That is, for the MSE between $f_N$ and $f_\\infty$, we have $\\left< (f_N-f_\\infty)^2 \\right> =  (\\left< f_N \\right> - f_\\infty)^2   + \\left<(f_N - \\left< f_N \\right>)^2 \\right>$.\n2. **'Feature learning in finite width SP networks'**: We completely agree that finite width SP networks display feature learning. Our aim was to argue that because infinite width SP networks do not have feature learning, infinite width SP networks cannot explain the dynamics of the finite width networks we train in practice which do exhibit feature learning. Therefore SP networks can not be consistent across all large widths. We will rewrite to make this clear. \n3. **'Usage of mup networks instead of SP networks'**: This is not true since some of the networks we look at are indeed SP network. Note that $\\mu$P should not be thought of as a different type of network than finite width standard/NTK parameterization networks, but rather as a rule for how to make a finite width network wider without sacrificing feature learning (preventing a limiting kernel behavior or unstable dynamics). Each $\\mu$P network is defined by a base network and a rescaling of the width. If the rescaling factor is set to 1 we recover the original network. We take our base networks to be SP networks (Appendix A.2), hence a) width 64 networks for the CIFAR-5m and ImageNet and b) width 256 for Wikitext are exactly SP networks. So in fact our results argue that SP networks have similar dynamics (up to the discussed deviations) as their corresponding infinite width $\\mu$P analogues. We will add a discussion to clarify this. To see an example showing that other parameterizations do not lead to consistent dynamics, please see Figure 1 in the attached rebuttal PDF which shows that kernel and predictor dynamics change with width in NTK parameterization. The base networks at width $N=128$ agree perfectly, but as the width $N$ is increased the two parameterizations ($\\mu$P and NTK) approach different limits.\n4. **‘How small of a difference on a metric like cross-entropy loss is needed in order to consider it a negligible difference, or “consistent” as defined in this paper?’**: To quantitatively study convergence with width, in Figure 3 of the rebuttal pdf we fit $C, D, \\beta$ in the equation $L = C (width)^{-\\beta} + D$ via the scipy curve_fit function. We find an excellent fit to the observed data for both language and vision tasks. This is consistent with diminishing returns of increasing width. We plotted this for a given training step near the end of training, but consistently observe the same power-law trends in $N$ throughout training.  \n5. **‘Ensembling in logit plots’**: Indeed we are averaging those plots over initialization. We will make sure to add this to the discussion along with error bars showing the standard deviation.\n6. **'Benign overfitting on test'**: Both bias ($(\\left< f_N \\right> - f_\\infty)^2$, as discussed above) and variance tending to 0 for test predictions implies that test predictions are converging with the widths used in the experiment. On the other hand this does not happen for train where larger widths overfit more i.e. more overfitting on train at larger widths does not lead to worse test performance. This is why this is an instance of benign overfitting. We will expand the discussion to make this clear.\n\n\nWe thank the reviewer for pointing out that our language was confusing in Section 4. We will make our explanations more concise and clear. Below we answer some of the questions \n    1. We will fix this sentence, writing instead that  \"... although the eigenvalue spectrum of the ensembled eNTK is not substantially affected by finite width, the correlation of the top eigenfunctions with the target function decreases with width.\"\n    2. The \"task\" is simply fitting the target polynomial $y(x)$ in the online setting. We will remove the use of \"task\" for $y(x)$ in the text.\n    3. We wanted to identify a setting where a simple explanation could account for the bias gaps across widths. For gradient flow, the bias correction has to do with changes in the statistics of the NTK $K$ across widths. In principle, the eigenvalues and eigenvectors of $K$ could change, but we found that the difference in the dynamics was largely due to changes in the eigenvectors. \n    4. We will use a simpler vernacular. Instead of referring to decompositions and \"spread\" of power, we will switch to using \"correlation\" language. For instance, the correlation of eigenfunction $k$ with the target is $\\left< y(x) \\psi_k(x) \\right>$. The test MSE depends on these correlations and on the eigenvalues $L = \\sum_{k} \\left< y(x) \\psi_k(x) \\right>^2 e^{-2 \\lambda_{k} t}$. We will then refer to $C(k)$ as the correlation of $y$ with the top $k$ subspace of eigenfunctions $C(k) = \\sum_{\\ell=1}^k \\left< \\psi_k(x) y(x)\\right>^2_x$. This metric is merely meant to capture the alignment of the top eigenvectors with the target function.\n\n\n**'figures/fonts in figures are far too small'**:  We will make sure to correct this by increasing the font and figure size (we will have one extra page for the final version).\n\n**'missing intuitions, explanations, contextualization'**:  We will make sure to do this in the final version (we will have one extra page for the final version). If there is a specific explanation that the reviewer finds missing we would be happy to provide it now."}, {"review_id": "lMZNdPtEcg", "reviewer": "Reviewer_j9zd", "summary": "This paper studies the convergence of key quantities (like test predictions, feature kernels) in Neural Networks as width increases, in feature-learning regimes. The paper is largely an empirical study, and presents width-convergence results on both image and language tasks. The paper also uses ensembling to consider the role of finite-width variance from randomness in initialisation, and finishes with a toy model that looks into the causes behind the bias in finite-width models compared to their infinite width counterparts.\n\n**Update: after clarifications in the rebuttal and additional experiments, I am raising my score to 6**", "questions": "Please respond to be weaknesses above, and also answer:\n\n1. Do you set the output weights to 0 in order to reduce the variance from function initialisation in your results? If not, then it might be worth looking at e.g. for the ensembling results.\n2. Do the experiments use identical mini-batch orderings across different widths?\n\nAlso, please correct the following typos:\n\n1. 'we put all of our in...' (line 37)\n2. 'the training lost' (figure 2 caption)", "limitations": "There isn't much discussion of limitations within the submission.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The question that this paper seeks to address, i.e. understanding finite-width NNs, is important and open.\n2. The paper has some interesting conclusions e.g. the convergence of feature-learnt test predictions at realistic widths (though see weaknesses too for some concerns), or the breakdown of finite-width bias in terms of eigenfunctions rather than eigenvalues.\n3. The paper is well-written and clear.", "weaknesses": "1. My main concern is that some of the results are not convincing/not matched to the claims of the paper. For example, for the qualitiative results on ImageNet experiments on Figure 2b and 3b, it isn't clear to me that these curves have converged in width (say comparing N=256 or 512). For the quantitative equivalent to 3b, Figure 3e, it's a bit surprising that the initial y-intercepts on the true logits RMSE with width 512 are so large, even at width N=256 (and seemingly not consistent with the claim that these predictions have converged at these evaluated widths). This is especially confusing given that it seems one of the claims in Figure 2 is that ImageNet results have converged for earlier training times compared to later times, and suggests that the reason why the True logit RMSE to width 512 is small at later epochs is simply because different networks are converging to the same correct prediction, rather than any width-limit occuring. Similar conclusions (larger RMSE at earlier training steps) appear to be the case in Figure 3f for WikiText-103.\n2. Likewise, I find the fact that we are often judging convergence in width off of 3-5 widths (e.g. Figure 1a-c) to be a bit unconvincing. I realise computational constraints are likely at play here, but it would be very good for the contribution of the paper to have more widths presented. Moreover, the argument that \"a quantity's deviations between successive widths decrease as widths are increased\" (line 54) is evidence for consistency is a bit weak to me: the harmonic series blows up to infinity despite satisfying this property.\n3. Going back to the motivation (e.g. line 86), the extent to which establishing convergence to infinite-width feature learning limits helps the community in understanding finite-width NNs might be quite modest, as such limits are very difficult to understand (significantly harder to understand compared to say kernel regimes), and also computationally prohibitive (e.g. Yang and Hu 2020, or Bordelon and Pehlevan 2022)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for these thoughtful comments. We'd like to address them in sequence. \n\n* *for the qualitiative results on ImageNet experiments on Figure 2b and 3b, it isn't clear to me that these curves have converged in width (say comparing N=256 or 512).*\n\nThis is a good point. In the main paper, we have not included enough plots of quantities as functions of $N$. This is something we plan to fix. In Fig 5 of our supplement we have plotted the final loss value as a function of $N$ (as well as ensemble number) for the Imagenet experiment in 2b. In Fig 3, A scipy.curve_fit on our losses as functions of $N$ gives an excellent fit to the power law $C N^{-\\beta} + D$.  A similar fit works for the transformer model as well. Moreover, the exponent $\\beta$ is frequently found to lie between 0.3 and 1.0. This power law dependence is also present in figures 4c, f, although is more cleanly seen when tracking losses rather than representation differences. \n\nFor Fig 3b, it is true that the logits tend to grow over the course of training. In this setting, we could have plotted the change relative  to the logit of the widest network, which would give a curve similar to 3e. For 3e we did in fact divide by the widest network's logit variance to rescale everything and keep things constant in time. We will explicitly note this in the updated revision. \n\n* *\"...it's a bit surprising that the initial y-intercepts on the true logits RMSE with width 512 are so large...\"*\n\nOne important aspect that we will clarify in the main text is that in calculating this RMSE we are dividing by the standard deviation of the widest network's outputs. That is:\n$RMSE = \\frac{\\sum_i (y_i^n - y_i^{widest})^2 }{\\sum_i (y_i^{widest})^2 }$\nwhere $i$ indexes the test points and $n$ is the width of the network we are comparing to. Near the start of training, the outputs are small, especially in $\\mu$P, and this causes an early sensitivity in the curves. Once the network's output has grown, it goes away. This is the primary driver of this effect.\n\n* *I find the fact that we are often judging convergence in width off of 3-5 widths (e.g. Figure 1a-c) to be a bit unconvincing.*\n\nIt would indeed be nice to go to still-higher values of $N$ but as you point out, we eventually do become limited by compute. For the transformer model, the width 4096 is quite large. Larger width models cannot easily be loaded onto an 80GB A100. We sweep over 7 widths and 2 orders of magnitude from width 64 to 4096 at good resolution. For the imagenet models, the ResNet architecture has $N$ for the initial number of channels. The final number is $8N$, or 4096 for the width 512 network. This is again near our computational limit, and is realistic for what is used in practice.. The range of widths was largely motivated by matching with the model sizes used in practice, to the extent that we could. Finally, the power law fits we observe match the data to high accuracy across the entire range of widths, providing some evidence that our observed trends can be extrapolated to a reasonable degree. \n\n* *Moreover, the argument that \"a quantity's deviations between successive widths decrease as widths are increased\" (line 54) is evidence for consistency is a bit weak to me: the harmonic series blows up to infinity despite satisfying this property.*\n\nThis is an apt observation, and one that worried us when formulating a definition of consistency. When we refer to consistency we simply state that with high probability a given network observable lies in an increasingly small region as the width grows, using language similar to that defining Cauchy sequences. In practice it is better to do a power law fit on a given network quantity and see if the fit is accurate and decaying. With a finite number of data-points this is likely the best one can do. All of our decays are well-fit by power laws with non-negligible exponent in $N$ (see above comment and Fig 3). A diverging sequence like the partial sums of the harmonic series would exhibit a power law exponent that either is close to zero or gets smaller as more points are added. \n\n\n* *\"Going back to the motivation...\"*\n\nIndeed, the simulations of mean field limits require a time cubic in the number of train steps to approximate. Even kernel methods are for the most part prohibitively expensive for large-scale datasets seen in practice, but mean field limits are so far less interpretable. We believe that a key contribution of this work is that, by contrast to NTK limits, which frequently break at realistic widths due to both finite-width kernel corrections and NTK changes due to feature learning, the $\\mu$P limits can be reached achieved at realistic widths to high accuracy at early training times. Further, there are still nontrivial predictions that can be made for infinite-width mean field networks and their finite width corrections even if simulating them is complicated.\n\n**Questions** \n\n* *\"Do you set the output weights to 0 in order to reduce the variance from function initialisation in your results? If not, then it might be worth looking at...\"*\n\nWe did not originally have the outputs subtracted off. We did however look at this change and observed no major qualitative difference in the results. We have included this result as figure 5c in the supplement and will incorporate it as an appendix figure. \n\nIn the lazy limit, this subtraction is particularly important, as the output at initialization is $O(1)$ in the width $N$. For $\\mu$P, the output scale goes like $1/N$, and we expect such a subtraction to be less important.\n    \n* *Do the experiments use identical mini-batch orderings across different widths?*\n\nYes, we do use the same mini-batch orderings across all widths. We will revise the text to state this explicitly.\n\n* *\"Also, please correct the following typos:...\"*\n    \nCertainly. We have corrected the typos pointed out in our current draft copy. Thank you for catching them. "}]

Author rebuttal:
We thank the reviewers for their detailed reading of our paper and useful suggestions and comments. In response to many repeated questions and concerns, we have performed the following new experiments (see attached PDF)

#### New Figures and Experiments
1. We have added a comparison of network dynamics for $\mu$P and NTK networks (without layernorm). While the $\mu$P networks quickly approach a consistent limit, the NTK networks exhibit different dynamics and train more slowly as width increases. Width consistency is harder to achieve in NTK parameterization (or standard param with learning rate $1/N$) since (a) kernels must concentrate and (b) feature learning effects must go to zero. Consistency in $\mu$P is a strictly weaker condition as it only requires (a) concentration. We will stress this point more in our our paper.
2.  We verify that layernorm does not significantly change our results, showing strong width consistency on CIFAR-5M. We also provide a theoretical argument that layernorm should not prevent a well defined infinite width limit, and argue that this falls under the Master Theorem of Yang & Hu 2021 (see response to reviewer ddkm). 
3. We plot the explicit scaling relationship for both Imagenet models and Transformer models at a fixed time, showing that a power law fits the data very well. 
4. We train a myrtle network without layernorm on CIFAR-5m and observe similar consistency in that setting as well. 
5. (5a) We provide plots that vary the number of ensembles to verify that we have used a large enough ensemble count (20) to estimate average loss and bias.
5.  (5b) Similarly, we show that returns to making network wider exhibit diminishing returns by plotting performance at a fixed time as a function of $N$.
5. (5c) We check whether zeroing out the last layer at initialization has a significant effect on the dynamics/consistency, and find little change.

#### Writing updates
In addition to the new experiments we also aim to make some parts of our paper more clear. Specifically, we hope to clarify the following points 
 1. We will provide more citations and connections between our work and existing empirical papers which study width consistent dynamics (see Reviewer ddkm). Some prior works are concerned with similar effects (eg fluctuations due to finite size). We think our paper is novel in that we have stress tested the infinite feature learning limit in experiments on larger scale networks and datasets and in the $\mu$ parameterization and examined many metrics including internal losses, representations, large step size effects, etc (see response to ddkm).
2. We will clarify what we mean by *bias* and *variance*. By bias, we merely refer to the deviation between the ensembled predictor and infinite width predictor $(\left< f_N \right> - f_\infty)^2$. For the MSE between $f_N$ and $f_\infty$, we have the typical bias variance decomposition $\left< (f_N-f_\infty)^2 \right> =  (\left< f_N \right> - f_\infty)^2   + \left<(f_N - \left< f_N \right>)^2 \right>$. We note that this is not the only type of bias/variance decomposition used in Machine Learning theory (Adlam & Pennington arXiv:2011.03321).
3. We will explain why we resort to experiments to try establishing convergence to infinite width rather than computing the infinite width theory. Computing the feature learning infinite width limit is challenging (super exponential time for exact numerical results, or $\sim T^3$ time for $T$ steps of SGD with Monte-Carlo methods). 
4. Since finite size effects do tend to accumulate in large scale settings and this is an important aspect of our empirical results, we change our sentence in the discussion from "our results motivate studying infinite width feature learning limits" to "our results motivate studying infinite width limits and the accumulating finite width deviations from this limit."
5. We will acknowledge our use of our widest networks as a proxy for the infinite width limit as a limitation of our results in the discussion. 
6. We will make the theoretical section of our paper more clear and readable and invoke minimal jargon. With the additional page of space, we can provide longer explanations of what we aim to establish there (explaining the source of bias gaps in a simple setting). We will use a more standard vernacular and avoid convoluted phrases about "spread of task power" etc. 
7. We will be clear about how we are averaging our plots over random seed and will add error bars for standard deviation over inits where applicable. We will also make the fonts and figures bigger so that they are more readable. 
8. We will clarify that $\mu$P should not be thought of as a different type of network than finite width standard/NTK parameterization networks, but rather as a rule for how to make a finite width SP network wider without sacrificing feature learning (preventing a limiting kernel behavior or unstable dynamics). See figure 1 in attached PDF for a clarification.   

Meta-review: The paper shows realistic networks achieve infinite-width feature-learning behavior throughout training for easy tasks and for early training time in hard tasks. The authors conducted an empirical investigation into the consistency across widths of the training dynamics of neural network models in the maximum-update scaling. Rather than using the neural tangent kernel as the descriptor of infinite-width behavior, the paper focuses on a different parameterization called feature-learning networks, which vary from standard-parameterization networks in that intermediate layer weights do not remain fixed even in the infinite-width limit. 

The reviewers pointed out this work provides strong, novel empirical study. The experimental  coverage is wide enough to extract practical applicable lessons. Some notable quotes from the reviewers:

- "the first to consider the empirical investigation of the large-width consistency of neural networks in the maximum-update parameterization" (`ddkm`)

- "a wide range of numerical experiments on various tasks with different neural network models that cover several relevant aspects of the model training behavior." (`gHF9`)

- "appreciate that this paper includes a mixture of convolutional and attention-based neural networks on a mixture of vision and language tasks, and studies finite and infinite-width agreement across loss, per-sample predictions, and network weight" (`gHF9`)

Also during rebuttal authors added empirical comparison of network dynamics for muP and NTK parameterized networks with regards to width consistency and further demonstrated weak dependence on inclusion of LN.

With growing interest in the feature learning limit of deep learning, broad empirical study on consistency of width in these limits is valuable for both practical applicability and further theoretical study in the future. As reviewers unanimously recommended accepting the paper (6, 6, 6, 6, 5), AC recommends acceptance.
Decision: Accept (poster)