Title: Feature-Learning Networks Are Consistent Across Widths At Realistic Scales
Venue: NeurIPS 2023
Accepted: yes
Keywords: mean-field

muP

feature learning

infinite width

deep ensembles
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

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.

1 Introduction

Studies of large-scale language and vision models have shown that models with a larger number of parameters achieve better performance [1, 2]. Motivated by the success of large-scale models, several theories of deep learning have been developed, including large-width limits. Infinite width limits which arise in standard parameterization (SP) or neural tangent parameterization (NTP) considered in [3, 4] gives rise to a a model with no feature learning (a kernel method). In this limit, the neural network loses the ability to adapt its internal features. Feature learning is crucial to explain deep learning's superior performance to kernels, the emergence of interpretable neurons such as edgedetecting CNN filters, transfer learning capabilities, and large learning rate effects such as edge of stability [5-8]. All of these effects are exhibited in modern large-scale networks.

Recently, several works have identified an alternative parameterization of neural networks that preserves feature-learning even at infinite width [9-14]. In this work, we focus on the maximal update parameterization ( µ P), or equivalently the mean field parameterization [12, 15, 14]. The existence of infinite-width feature-learning limit, suggests that this parameterization is potentially more promising to explain deep learning phenomena than the previous limits. This motivates us to ask:

Question: Can realistic-width neural networks be accurately described by their infinite-width feature-learning limits?

10

100

1000

10000

8

16

32

64

N=64

N=128

N=256

N=512

(a) CIFAR-5m

150

300

600

1200

2400

4800

Train Steps

4

5

6

7

Loss

384

768

1536

(b) C4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Images seen

1e6

0

20

40

60

80

100

True logit value on test point

N=32

N=64

N=128

N=256

N=512

(c) ImageNet

N = 128

N = 256

N = 512

N = 1024

N = 2048

N = 4096

(d) Wikitext-103

We attempt to answer this question by training networks of varying widths on vision and language tasks for realistic datasets and architectures. Concretely, we focus on on the online setting, where data is not repeated during SGD, and track the following quantities across widths:

The losses throughout training.

The predictions of the networks on individual points throughout training.

The learned representations, summarized by the feature kernels preactivation distributions and, for transformers, their attention matrices.

Dynamical phenomena such as the edge of stability governing the top Hessian eigenvalues, as well as large learning rate and small batch size effects on the loss.

On each of these metrics, we show that sufficiently wide neural networks converge to consistent behavior across widths. In Figure 1, we show loss curves, logit predictions, and attention matrices approach consistent behavior as width is increased across several architectures and datasets. We further observe that the widths that achieve this consistent behavior are within the range of those used in practice. We use large-width consistency as a proxy for achieving the limiting infinite-width behavior. We stress that this observed consistency is a property of networks in mean field/ µ P parameterization but is not present in other parameterizations which also give an infinite width limit like NTK parameterization (See Appendix D for a comparison).

We say that a network property is consistent if, beyond some width, its values all lie within some small interval with high probability. We measure consistency by showing that a quantity's deviations between successive widths decrease as the widths are increased, and that its value for narrower networks systematically approaches its value for the largest trained network.

Our results show the following:

For simple vision tasks such as CIFAR-5m [17], ResNets with practical widths achieve near consistent loss curves across widths (Section 2).

Beyond the loss curves, the individual predictions of the networks agree pointwise. That is, the logits agree on test points throughout the training process. We further show that internal representations as measured by distributions of neuron preactivations and feature kernels in various layers are consistent across widths (Section 2).

For harder tasks such as ImageNet and language modeling, loss curves are consistent across widths early in training. As training progresses, loss curves for narrow networks deviate smoothly from the loss curves of wider networks. The effective width required to reach infinite-width behavior thus increases with training time. Conversely, as network size grows we approximate the infinite width network for a larger number of training steps (Section 2).

Finite-width neural networks have variance in the learned function due to initialization seed. This variance depends inversely on the width. We study ensembles of networks over different initializations to remove this noise. Further, by training ensembles of networks, we can perform

a bias-variance decomposition over initializations (c.f. Appendix F for details and definitions) to analyze the effects of finite width. We find that finite-width bias plays an important role. Equivalently, ensembling narrow networks does not yield infinite-width behavior (Section 3).

In the setting of offline learning, at late times one can over-fit the training set. We observe that this leads to larger gaps in network behavior across widths, and can break the trend that wider networks perform better (Section 3).

We develop a spectral perspective on the origin of the finite-width bias by analyzing it in a simple setting of a lazy network learning a simple task. We then apply this perspective to a CNN trained on CIFAR-5m (Section 4).

The consistency across large widths strongly suggests that the dynamics and predictions of realisticscale networks can be effectively captured by their infinite-width feature learning limits. For realistic tasks, as the width is increased, a larger interval of training can be characterized by this infinite-width limit. Most importantly, even though quantitative agreement across widths slowly breaks with more and more training, we observe that wider networks perform better (as in [15]) and preserve qualitative aspects of the learned features (such as hidden layer kernels and attention matrices) and dynamical phenomena such as edge of stability. This suggests that infinite width feature-learning networks are good models to study deep learning.

Our results have implications for interpretability, as the agreement of internal representations suggest that many other phenomena, such as transfer learning with linear probes or fine-tuning, in-context learning [18, 19], the emergence of outliers [20], and the emergence of induction heads [21] may be understood from the perspective of infinite-width feature learning networks.

We plan to have our code made freely available on github to ensure the reproducibility of these results.

1.1 Related Works

Empirically, the scaling of relevant quantities with width in the standard or neural-tangent parameterizations was thoroughly studied in [22]. In the latter parameterization, sufficiently wide networks give a kernel method with the infinite-width NTK. Several papers have shown that in practice the NTK limit insufficiently characterizes realistic deep neural networks [5, 23, 6]. Attempts to capture feature learning and predictor variance from perturbative series around infinite-width dynamics show that finite-width variance and kernel adaptation scale as 1 /N [24-26] for width N . A 1 /N scaling of generalization error with width was empirically verified on many tasks [27, 28]. The effect of width on generalization in the feature-learning regime was empirically studied in [29] in the relatively limited setting of multi-layer perceptrons (MLPs) on polynomial tasks. There, the variance of the finite-width NTK at the end of training adversely affected generalization. Bias-variance decompositions over dataset, label noise, and initialization parameters were studied in [30] for linear models.

The authors of [31] identified that altering the output scale α of any network could increase or decrease feature learning in a neural network. Large values of α correspond to the 'lazy limit' where the network's features don't evolve. A follow up study noticed that rescaling the output by α = α 0 / √ N for width N networks gave consistent behavior of feature learning and losses in small scale experiments [10]. Several works have studied this regime of training in the two-layer limit, known as 'mean field' parameterization, where features are still learned even at infinite width [32, 9, 33, 34]. Extensions of this model to deeper networks were studied in [35-38, 12, 14]. A theory of finite-width corrections to networks in this parameterization was studied in [39]. A very general set of parameterization principles, termed µ P, was introduced to give a well defined feature learning limit for a wide range of architectures including RNNs, CNNs, MLPs and transformers [12]. [15] demonstrated that this parameterization nearly fixes optimal hyperparameters across network widths, allowing for hyperparameter transfer from small to large widths. This work also empirically noted that wider networks always outperformed narrower networks in this parameterization.

Our paper focuses on networks in µ P and attempts to study the consistency of many relevant network properties across widths. We perform a fine-grained analyses of more realistic models throughout the dynamics of training. To the best of our knowledge, this is the first such paper to study the consistency of network outputs, internal representations, and dynamics across widths.

10 1

10 2

10 3

10 4

Train Steps

0.5

1.0

1.5

2.0

Training Loss

N

=32

N

=64

N

=128

N

=256

N

=512

(a) CIFAR-5m

10 5

10 6

Images seen

2

3

4

5

6

7

8

Training Loss

N=32

N=64

N=128

N=256

N=512

(b) Imagenet

(c) Wikitext-103

5000

10000

15000

20000

Train Steps

0

1

2

3

4

5

6

7

True logit value on test point

N=64

N=128

N=256

N=512

(a) CIFAR-5m

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Images seen

1e6

0

20

40

60

80

100

True logit value on test point

N=32

N=64

N=128

N=256

N=512

(b) Imagenet

0.0

0.2

0.4

0.6

0.8

1.0

Tokens Seen

1e8

5

6

7

8

9

10

11

True logit value on test point

N = 64

N = 128

N = 256

N = 512

N = 1024

N = 2048

N = 4096

(c) Wikitext-103

5000

10000

15000

20000

Train Steps

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

True logit RMSE from N = 512

32

64

128

256

(d) CIFAR-5m

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Images seen

1e6

0.0

0.2

0.4

0.6

0.8

1.0

True logit RMSE from

N

=512

N=32

N=64

N=128

N=256

(e) Imagenet

0.0

0.2

0.4

0.6

0.8

1.0

Tokens Seen

1e8

0.10

0.15

0.20

0.25

0.30

0.35

True logit RMSE from

N

= 4096

N = 64

N = 128

N = 256

N = 512

N = 1024

N = 2048

(f) Wikitext-103

2 Consistency of large-width behavior in online learning

We focus on studying the effect of width in the setting of neural networks learning a task in the online setting. Online learning is representative of many modern settings of deep learning, and as will be shown in Section 3, obviates consideration of memorization and over-fitting in offline learning that can lead to large differences in networks across widths.

In what follows, the variable N will denote the width of a given network. For vision tasks, this will correspond to the number of channels in each layer. For transformers, in the notation of [40], N = d model = hd k = hd v and d ffn = 4 N . Here, h is the number of heads, which we will keep

1.0

5000

10000

15000

20000

Train Steps

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

True logit RMSE between widths

(32, 64)

(64, 128)

(128, 256)

(256, 512)

(a) CIFAR-5m

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Images seen

1e6

0.0

0.2

0.4

0.6

0.8

True logit RMSE between widths

(32, 64)

(64, 128)

(128, 256)

(256, 512)

(b) Imagenet

0.0

0.2

0.4

0.6

0.8

1.0

Tokens Seen

1e8

0.08

0.10

0.12

0.14

0.16

0.18

0.20

True logit RMSE between widths

(64,128)

(128,256)

(256,512)

(512,1024)

(1024,2048)

(2048,4096)

(c) Wikitext-103

fixed. d model is the embedding dimension of the tokens as well as the dimension of the residual stream. d k is the dimension over which the dot products in the attention are calculated and d v is the dimension of the values in the attention layers. d ffn is the hidden width of the feedforward networks (FFN).

Convergence of loss curves We begin by showing (Fig. 2) that the loss curves for sufficiently wide networks on a given task achieve consistent behavior across widths. Throughout the paper we measure train loss in terms of crossentropy. For all tasks, at early times large widths agree, but for more complicated tasks such as ImageNet or Wikitext-103, learning curves of narrower network deviate from those of wider ones.

The width beyond which networks emulate infinite-width behavior depends on the complexity of the task. For more difficult tasks, larger widths are required for the loss curves to converge. For simple tasks such as CIFAR-5m we find that widths as narrow as 128 are essentially consistent with infinite width-behavior for an entire pass through the 5 million image dataset. For ImageNet, widths near 512 are close to consistent for four passes through the dataset with heavy data augmentation. These widths are well within the range of those practically for images [41, 42]. For transformers going through a single full pass of Wikitext-103, widths on the order of 4000 are required. Early transformer models certainly had hidden widths of order 4k [43], and more recent models such as GPT-3 have widths going up to 12288 [18], so this is also within the regime of realistic width.

Pointwise convergence of predictions Beyond the convergence of the training loss curves, we observe that the logits of a network on a fixed test point become consistent as width grows. This test point can be an image in the test set or a masked token in the validation set. In plots a), b), and c) of Figure 3, we show that for a specific held-out test point, the value of the network on the correct logit becomes consistent as the width grows. In d), e), and f) we plot the root mean squared distance to the widest networks logits over the test set. We further study the difference between successive widths in Figure 4.

Convergence of representations In addition to loss and prediction dynamics, we also examine whether learned representations in these models are consistent across widths. Mean field theories of neural network dynamics predict that sufficiently wide networks should have identical kernels (and attention matrices for transformers) and that all neurons in a layer behave as independent draws from an initialization-independent single-site distribution [9, 12, 14, 46, 47]. To test whether realistic finite-width feature learning networks are accurately captured by this limit, in Figure 5, we analyze the feature kernels and preactivation distributions before and after training as well as the attention matrices in transformer models trained on Wikitext-103. We see qualitative consistency in the plots of kernels and attention matrices in b) and c) which can be made quantitatively precise by plotting the distance to the widest networks and showing systematic convergence in c) and f).

Convergence of dynamical phenomena In Figure 6a, we show that the sharpness, defined as the top eigenvalue of the loss Hessian, grows steadily to a final value that it then fluctuates around. This is a small-batch analog of the the edge-of-stability phenomenon identified in [7]. We also show in Figure 6b that on CIFAR-5m task, at early times, the individual variations due to batch noise and large learning rate effects can be consistently captured across widths for µ P networks. In Appendix

5

0

5

h

10

3

10

2

10

1

p

(

h

)

Initial

N = 64

N = 128

N = 256

N = 512

N(0.0,3.0)

5

0

5

h

10

3

10

2

10

1

Final

N(-0.8,1.6)

(a) Preactivations CIFAR-5M

Initial

N

=64

Final

N

=128

N

=256

N

=512

(b) Last Layer Kernels

32

64

128

256

N

10

2

10

1

1-CKA(

4

N

,

4

512

)

N

1

N

1/2

=1

=2

=3

=4

(c) Convergence of Kernels

10

5

0

5

10

Activation

10

4

10

3

10

2

10

1

Density

Layer 4 FFN, before training

N=64

N=256

N=4096

N(0.0, 1.0)

10

5

0

5

10

Activation

10

4

10

3

10

2

10

1

Layer 4 FFN, after training

N(-1.1, 2.2)

(d) FFN Preactivations Wikitext

N = 128

N = 256

N = 512

N = 1024

N = 2048

N = 4096

(e) Attention Matrices

64

128

256

512

1024

2048

N

10

2

10

1

KL Divergence from width 4096 model

N

1

N

1/2

DKL

(

A

4096

|

A

), =1

DKL

(

A

4096

|

A

), =2

DKL

(

A

4096

|

A

), =3

DKL

(

A

4096

|

A

), =4

(f) Convergence of Attention

1

10

100

1000

Train Steps

0

50

100

150

200

250

Top Eigenvalue

64

128

256

(a) Edge of Stability

0

10

20

30

40

50

60

Train Steps

1.9

2.0

2.1

2.2

2.3

2.4

Train Loss

64

128

256

E, we further demonstrate sharp agreement of large learning rate and small batch size phenomena for MLPs learning a simple task. There, we show that while µ P leads to strikingly consistent loss curves, SP does not.

3 Deviations from large-width behavior

The consistency observed in Section 2 may break later during training in either the online or offline settings. In the online setting, deviations owing to narrow width compound over time and lead to two sources of error relative to the infinite width limit which we describe in 3.1. In the offline setting, where data is recycled several times, networks over-fit the training data, which can lead to larger gaps between widths and can break the trend that wider networks perform better.

Finite-width effects introduce an initialization dependence to the network, leading to additional variance in the learned function and hindering generalization [27-29]. This initialization-dependent variance can be mitigated by averaging the output logits of a sufficiently large ensemble of networks [48]. Using the bias-variance decomposition terminology, we refer to the discrepancy in performance between an ensembled network and the expected performance of a single network the variance , and the gap between an ensembled network and the behavior of infinite-width network as the bias of narrower width. We elaborate thoroughly on what we mean by this decomposition in Appendix F. By definition, the expected difference in loss between a single finite-width network and an infinite-width network is the sum of the bias and the variance. Below, we investigate the behavior of bias and variance in networks across various vision and language tasks.

3.1 Online training

10 1

10 2

10 3

10 4

Train Steps

0.25

0.5

1.0

2.0

4.0

Training Loss

Single Network

Ensemble

N = 32

N = 128

N = 512

(a) CIFAR-5m

10 5

10 6

Images seen

2

3

4

5

6

7

Training loss

N=32

N=128

N=512

Single Network

Ensemble

(b) Imagenet

10 6

10 7

10 8

Tokens Seen

4.0

4.5

5.0

5.5

6.0

6.5

7.0

Loss

Single Network

Ensembled

N = 64

N = 256

N = 1024

N = 4096

(c) Wikitext-103

Figure 7 shows that at large widths, both single networks and ensembles of networks achieve comparable error. In this regime, all the networks are consistent and increasing the width has a very marginal effect, as does ensembling. At narrower widths, variance is nontrivial (i.e. ensembling helps) but bias is much larger than variance. Single wide networks outperform ensembles of narrower networks. By comparing a) with b) and c) of Figure 7, we see that harder tasks induce larger bias gaps. Prior theoretical work [28, 29] has focused mostly on studying the variance term. In Section 4 we study the bias from a theoretical perspective.

3.2 Offline Training

In offline learning, which refers to multi-epoch training, we encounter several unexpected phenomena that challenge the width consistency observed in the previous section, even at large widths. To compare offline learning with online learning, we utilize CIFAR-200k, a 200k sized random subset of CIFAR-5m. Previous studies have demonstrated that label noise contributes to an increase in overfitting [49]. In order to investigate how width consistency changes with overfitting and double descent, we conduct experiments on a noisy label version of CIFAR-50k (50k sample from CIFAR5m), where 50% of the labels are noisy. Additional ImageNet experiments are presented in Appendix J. As offline training achieves near-zero error, we need to compare very small quantities. To accomplish this, we will plot and compare all quantities on a logarithmic scale. The following phenomena are observed:

Single network performance on the training set does not converge with width, even at high widths (Figure 8 (a)). In other words, the combined bias and variance does not reach zero, even with substantial widths. This is in contrast to the online runs.

Ensembling (Figure 8 (b)) reveals that both bias and variance terms individually fail to reach zero, even at high widths.

Regarding test performance, both bias and variance tend to zero as width increases, demonstrating an instance of benign overfitting (Figure 8 (d) and (e)).

When working with the noisy label version of CIFAR-50k, we observe clear overfitting and stepwise double descent [49] as training progresses (Figure 8 (f)). Notably, we observe significant

10

100

1000

10000

Train Steps

1

4

16

64

Train Error

N=128, Offline

N=128, Online

N=256, Offline

N=256, Online

N=512, Offline

N=512, Online

(a) CIFAR-200k, offline vs. online

10

100

1000

10000

Train Steps

1

4

16

64

Train Error

N=128, Offline

N=128, Offline Ensemble

N=256, Offline

N=256, Offline Ensemble

N=512, Offline

N=512, Offline Ensemble

(b) CIFAR-200k, offline ensembled

1000

2000

4000

8000

Train Steps

0.001

0.01

0.1

1

10

100

Train Error

64 avg error

64 ensemble error

256 avg error

256 ensemble error

10

100

1000

10000

Train Steps

4

16

64

Test Error

N=128, Offline

N=128, Online

N=256, Offline

N=256, Online

N=512, Offline

N=512, Online

10

100

1000

10000

Train Steps

4

16

64

Test Error

N=128, Offline

N=128, Offline Ensemble

N=256, Offline

N=256, Offline Ensemble

N=512, Offline

N=512, Offline Ensemble

1000

2000

4000

8000

Train Steps

20

40

60

80

Test Error

64 avg error

64 ensemble error

256 avg error

256 ensemble error

(d) CIFAR-200k, offline vs. online.

(e) CIFAR-200k, offline ensembled

Figure 8: Top Row: Effects of offline training on train metrics. In (a) and (b) we do multi-epoch training on CIFAR-200k. We see that both bias and variance for train error are magnified by offline training and do not tend to 0 for the largest widths we could try. (c) we do multi-epoch training on noisy CIFAR-50k and again observe large bias and variance terms at large widths. Bottom Row: Effects of offline training on test metrics. In (d) and (e) we do multi-epoch training on CIFAR-200k. We see that both bias and variance for test error are near 0 at large widths. (f) We train on noisy CIFAR-50k and observe that 'wider is better' is violated for ensembled networks.

deviations in width for single network performance, indicating that the benign overfitting observed in Figure 8 (d) and (e) is dataset-dependent. Furthermore, variance is found to be much larger than in the non-noisy experiments.

Surprisingly, we discover (Figure 8 (f)) that some ensembled narrower width networks outperform ensembled wider networks. This presents a counterexample to the 'wider is better' phenomenon [15] for ensembled networks. We hypothesize that such counterexamples can only exist in the context of offline training.

4 Spectral perspective on finite-width bias

In this last section, we develop a toy model in which the effect of finite-width bias can be clearly seen. We analyze it first in the simple setting of an MLP fitting a polynomial in the lazy limit. Here, all the dynamics are well-captured by the finite-width empirical neural tangent kernel (eNTK). By studying the spectral properties of this kernel across widths, we see that finite widths lead to eNTK's with worse bias components in their losses.

Concretely, we see that although the eigenvalue spectrum of the ensembled eNTK is not substantially affected by finite width, the decomposition of the task into eNTK eigenvectors changes, with narrowwidth eNTK's putting more of the task into smaller eigenmodes that take longer to be learned. In practice, applying this analysis to the after-kernel of the trained ResNets on CIFAR-5m reveals similar behavior. Prior literature has demonstrated that many of the properties of the final learned function are captured by the after-kernel [50-52].

We consider a model of online learning where a large batch of data from the population distribution p ( x ) is sampled at each step. This leads to approximate gradient flow dynamics d dt θ = -1 2 ∇ θ E x ( f ( x , θ ) -y ( x )) 2 (Appendix H). To analyze this equation, we choose a fixed orthonormal basis { ψ k ( x ) } for the space L 2 ( R D , p ( x ) d x ) of square-integrable functions on input space. The function f ( x ) , residual error ∆( x ) = y ( x ) -f ( x ) , and the empirical NTK

K ( x , x ′ , t ) can be expressed in this basis as f ( x , t ) = ∑ k f k ψ k ( x ) , ∆( x , t ) = ∑ k ∆ k ψ k ( x ) , and K ( x , x ′ , t ) = ∑ kℓ K kℓ ( t ) ψ k ( x ) ψ ℓ ( x ′ ) , respectively. Their training evolution is given by:

The statistics of the dynamical NTK matrix K kl ( t ) summarizes the statistics of the error dynamics ∆( x , t ) at any level of feature learning. At infinite width, K kℓ ( t ) is deterministic, while at finite width, it receives a Θ( N -1 ) mean displacement and a Θ ( N -1 / 2 ) fluctuation around its mean [24, 25, 39]. We consider approximating the dynamics of the ensembled predictor by d dt ⟨ f k ( t ) ⟩ θ 0 ≈ ∑ ℓ ⟨ K kℓ ( t ) ⟩ ⟨ ∆ ℓ ( t ) ⟩ . Here, ⟨·⟩ denotes averages over initializations. This expression neglects the contribution from Cov ( K kℓ , ∆ ℓ ) . We show that this approximation is accurate in depth-3 MLPs trained on Gegenbauer polynomial regression tasks in Figure 9 a). For more details see Appendix A.

In the lazy limit, the kernel is static and we choose ψ k to diagonalize ⟨ K kℓ ⟩ = δ kℓ λ k . This yields the loss dynamics L ( t ) = ∑ k ⟨ y ( x ) ψ k ( x ) ⟩ 2 e -2 λ k t . We can therefore quantify alignment of eigenfunctions to task with the cumulative power distribution C ( k ) = ∑ ℓ<k ⟨ y ( x ) ψ ℓ ( x ) ⟩ 2 x / 〈 y ( x ) 2 〉 x , which is the proportion of the task that is explained by the first k eigenvectors [53]. If C ( k ) rises rapidly with k then the loss falls faster [53]. In this limit, there are two ingredients that could make the bias dynamics across widths distinct. First, the eigenvalues λ k which set the timescales could be width-dependent. Second, the eigenfunctions ψ k ( x ) that diagonalize ⟨ K ⟩ can change with width. In Figures 9 b) and c) we show that the dominant effect is the latter. Finite-width corrections do not substantially affect the spectrum, but they do increase the proportion of the target function that lies in the eigenspaces corresponding to modes that are slower to learn.

To test whether these findings continue to hold in more realistic experiments, we computed the final NTKs (after kernels) of the ResNet-18 models trained on CIFAR-5M (specifically the models from Figures 3, 5). We ensemble average to get kernel ⟨ K c,c ′ ( x , x ′ ) ⟩ for output channels c, c ′ and input images x , x ′ . We then compute the kernel gradient flow corresponding to MSE training on the true target function for CIFAR-5M d dt ∆ c ( x ) = -∑ c ′ E x ′ ⟨ K c,c ′ ( x , x ′ ) ⟩ ∆ c ′ ( x ′ ) from initial condition given by the one-hot target labels ∆ c ( x ) | t =0 = y c ( x ) . The convergence rate of this dynamical system is again set by the eigenvalues and eigenfunction-task alignment. In Figure 9 (d), we find that the after kernels for wider networks give slightly more rapid convergence. Figures 9 (e) and (f) show that, similar to the MLP experiment, the spectra are very consistent across widths, but the eigenfunction task alignments, measured with C ( k ) are not. Overall, these experiments suggest that an important aspect of the bias of finite width models compared to their infinite width analogs is the deformation of their eigenfunctions.

5 Conclusion

We have demonstrated a striking consistency across widths for many quantities of interest to deep learning practitioners. Our fine-grained studies go beyond simply comparing test losses and have demonstrated that learned network functions, internal representations, and dynamical large learning rate phenomena agree for sufficiently large widths on a variety of tasks across vision and language. At later training times, or after many repetitions of the dataset, we observe systematic deviations brought on by finite width, and have characterized them in terms of the bias and variance of the network over initializations. This study motivates the applicability of infinite-width feature-learning models (and the accumulating finite width deviations from this limit) in reasoning about large scale models trained on real-world data.

In light of the accumulation of finite-width deviations at later training times, we caution that our study only exhibits the consistency of the infinite-width limit with the training time held fixed . It does not make a claim about other possibly limits that vary the training time jointly with the width to infinity, perhaps along a compute frontier. We leave further inquiry into such limits for future work.

0

50

100

150

200

250

300

350

400

t

10

4

10

3

10

2

10

1

10 0

tr

Losses for online MLP, L=3 k=2, batch size=250

NN

NN ens

K

N=499

N=1118

N=2499

(a) MLP Online training losses

10 0

10 1

10 2

10 3

10 4

Eigenvector Index

10 2

10 3

10 4

10 5

10 6

10 7

10 8

i

i

, MLP, L=3, k=2

N=100

N=223

N=499

N=1118

N=2499

10 0

10 1

10 2

10 3

10 4

Eigenvector Index

10

7

10

6

10

5

10

4

10

3

10

2

10

1

10 0

1

C

(

k

)

1

C

(

k

), MLP, L=3, k=2

N=100

N=223

N=499

N=1118

N=2499

(c) MLP unexplained task fraction

10 0

10 1

10 2

10 3

t

10

2

10

1

10 0

(

t

)

K

N = 32

N = 64

N = 128

N = 256

(d) CIFAR-5M After-Kernel

10 0

10 1

10 2

10 3

k

10

2

10

1

10 0

10 1

10 2

k

K

Spectra

(e) After-Kernel Spectra

10 0

10 1

10 2

10 3

k

0.0

0.2

0.4

0.6

0.8

1.0

1

C

(

k

)

K

Task Alignment

(f)

After-Kernel

unexplained task

fraction

Acknowledgments and Disclosure of Funding

We thank Boaz Barak, Jeremy Cohen, Alex Damian, Nikhil Ghosh, Gal Kaplun, Eric Michaud, Jamie Simon and Jacob Zavatone-Veth for helpful discussions throughout this project. We also thank Jacob Zavatone-Veth for comments on the draft.

AA is supported by the Professor Yaser S. Abu-Mostafa Fellowship from the Fannie and John Hertz Foundation. NV and DM are supported by a Simons Investigator Fellowship, DARPA grant W911NF2010021,and DOE grant DE-SC0022199. DM is supported by funding from the Office of Naval Research under award N00014-22-1-2377 and the National Science Foundation Grant under award CCF-2212841. SS is supported by a Susan Wojcicki and Dennis Troper Graduate Fellowship. BB is supported by a Google PhD fellowship. NV, BB, DM and CP are supported by funding from NSF grant DMS-2134157. CP is also supported by NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Compute was provided by the Harvard FASRC cluster and the Kempner Institute.

References

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.08361 , 2020.

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.15556 , 2022.

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.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems , 32, 2019.

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, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and HsuanTien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 405075699f065e43581f27d67bb68478-Abstract.html .

Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the ntk for understanding generalization in deep learning. arXiv preprint arXiv:2206.10012 , 2022.

Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065 , 2021.

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://distill.pub/2020/circuits.

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.

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 , 2020(11):113301, 2020.

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.

Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning , pages 11727-11737. PMLR, 2021.

Greg Yang and Edward Hu. Maximal Update Parametrization ( µ P) and Hyperparameter Transfer ( µ Transfer). https://github.com/microsoft/mup , 2022.

Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. arXiv preprint arXiv:2205.09653 , 2022.

Greg Yang, Edward J 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 https: //openreview.net/forum?id=Bx6qKuBM2AD .

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-5551, 2020.

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.

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-1901, 2020.

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.07682 , 2022.

Ziyang Luo, Artur Kulmizev, and Xiaoxi Mao. Positional artefacts propagate through masked language model embeddings. arXiv preprint arXiv:2011.04393 , 2020.

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.11895 , 2022.

Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems , 33:15156-15172, 2020.

Guillermo Ortiz-Jiménez, Seyed-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.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint arXiv:1909.05989 , 2019.

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2020.

Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory . Cambridge University Press Cambridge, MA, USA, 2022.

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):023401, 2020.

Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701 , 2021.

Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. arXiv preprint arXiv:2212.12147 , 2022.

Ben Adlam and Jeffrey Pennington. Understanding double descent requires a fine-grained biasvariance decomposition. Advances in neural information processing systems , 33:11022-11032, 2020.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems , 32, 2019.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics , 75(9):1889-1935, 2022.

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.

Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.00193 , 2019.

Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv:1902.02880 , 2019.

Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001.11443 , 2020.

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.

Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. arXiv preprint arXiv:2304.03408 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016 . British Machine Vision Association, 2016.

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.

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 arXiv:1909.08053 , 2019.

Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research , 13(1):795-828, 2012.

Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International Conference on Machine Learning , pages 3519-3529. PMLR, 2019.

Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications , 14(1):908, 2023.

Adam X Yang, Maxime Robeyns, Edward Milsom, Nandi Schoots, and Laurence Aitchison. A theory of representation learning in deep neural networks gives a deep generalisation of kernel methods. arXiv preprint arXiv:2108.13097 , 2021.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30, 2017.

Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020. URL https://openreview.net/forum?id=B1g5sA4twr .

Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105.10585 , 2021.

Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=1NvflqAdoom .

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186 , 2023.

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.

Zhengdao Chen, Grant Rotskoff, Joan Bruna, and Eric Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. Advances in Neural Information Processing Systems , 33:22217-22230, 2020.

A Experimental Details

A.1 MLPs

In Figure 9 we used a 3-layer MLP learning a Gegenbauer polynomial Q 2 ( β · x ) in D = 5 dimensions. Here β was a randomly chosen unit vector in R D . We implemented µ P parameterization by hand. The output layer of the network was rescaled by α 0 / √ N , consistent with µ P. We chose α 0 = 1000 to put us in the lazy regime. We set the the learning rate to be 5 N/ (1 + α 2 0 ) . General arguments based on kernel scale indicate that the learning rate should be scaled as α -2 0 at large α .

In Figure 14 we used a 3-layer MLP learning a Gegenbauer polynomial Q 2 ( β · x ) in D = 25 dimensions. We set the learning rate to be nearly as high as possible before a loss explosion.

A.2 Vision

A.2.1 CIFAR-5m

All plots except Figure 6a and 6b: We trained with standard CIFAR data augmentation of random crop (RandomCrop(32, padding=4) in pytorch) and horizontal flip (RandomHorizontalFlip() in pytorch). As base network (for µP ) we used ResNet18 where BatchNorm was replaced with LayerNorm (to maintain the consistency of the neural network between train and test). We used the SGD optimizer with learning rate of .05 with cosine decay over 20000 steps, .9 momentum and batch size of 250.

For Figure 6a, we used the above setup, but with a learning rate of 0.01 and a much higher batch size of 2000, so as to replicate the edge of stability phenomenon [7] which only occurs at high batch sizes. For Figure 6b, we used a learning rate of 0.3 and batch size of 32, so as to show the behavior of high learning rate and small batch size on train loss.

A.2.2 CIFAR-10 Multiple Passes

In Figure 10, we show the dynamics and representational consistency of ResNets trained on CIFAR10 for several epochs. The architecture is a ResNet18 with base-shape width set at N = 64 channels. The model is trained with SGD with learning rate 0 . 1 and cosine annealing schedule. The batch-size used is 128 .

A.2.3 ImageNet

In all ImageNet experiments, we used a training subset of the ImageNet-1k dataset consisting of 2 20 = 1048576 labeled images and a test subset consisting of 1024 labeled images. Both subsets were randomly sampled from the full ImageNet-1k training and validation datasets, respectively. To extend the duration in training in which the network remains in the online regime beyond one epoch, we heavily augmented the images in the training dataset using PyTorch's AutoAugment transform with the default policy, AutoAugmentPolicy.IMAGENET .

We again used the ResNet-18 architecture with µ P parameterization relative to the ResNet-18 network with base-shape width N = 64 channel [13]. All architectures and training procedures were implemented in Jax and used the auxiliary Flax and Optax packages, respectively.

Figures 2(b) and 7(b) were trained using the Adam optimizer with the following learning rate schedule: linear warm-up for 0.5 epochs from learning rate 8 × 10 -5 to 8 × 10 -3 , followed by cosine decay over 49 . 5 epochs to 8 × 10 -5 .

A.3 Language

A.3.1 Wikitext-103 Language Modeling

For all Wikitext-103 tasks, we adopted the µ P transformer as defined in the µ P package [13]. In the plots shown in the main text, we used a depth4 transformer, with d model , d k , d v = N and d ffn = 4 N . We performed a single pass through the train set in order to stay in the realistic online regime. We used a masked language modeling with sequence length S at varying input sequence lengths S . For Figure 1 d) we used the S × S attention matrix of an S = 128 transformer. In Figure 5 e) we used the attention matrix of an S = 35 transformer. We chose this different length simply to

illustrate the consistent message across sequence lengths. We used a batch size of B = 32 for all experiments. The residual stream was thus a tensor of shape ( S, B, d model ) .

We used the Adam optimizer with a learning rate of 0 . 0001 . We also ran the same configuration with SGD and a learning rate of 0 . 5 and observed the same behavior. See section B for further plots and details.

For figure 3, we used the Wikitext-103 validation set in order to measure the evolution of the predictions on masked logits. In 3 f), we averaged the mean squared error from the widest transformer by using 100 test points.

A.3.2 C4 Language Modelling

Figure 1 (b) we trained with base network being a 125m parameter transformer model on 2.5 billion tokens using the Mosaic ML's LLM codebase ( https://web.archive.org/web/ 20230519184343/https://github.com/mosaicml/examples/tree/main/examples/llm ). See https://web.archive.org/web/20230519183813/https://github.com/mosaicml/ examples/blob/main/examples/llm/yamls/mosaic_gpt/125m.yaml for the full hyperparameter details. We were limited by time and computational resources in our ability to explore further details of the C4 transformer model.

B Further Plots of Convergence

0

5

10

15

20

Epochs

10 1

10 2

Test Error

N = 32

N = 64

N = 128

N = 256

N = 512

(a) Error Consistency

6

4

2

0

2

4

6

h

10

4

10

3

10

2

10

1

10 0

p

(

h

)

Initial

N = 32

N = 64

N = 128

N = 256

N = 512

2

1

0

1

2

h

10

3

10

2

10

1

10 0

Final (20 epochs)

(b) Last Layer Preactivation Consistency

Initial

N = 32

Final (20 epoch)

N = 64

N = 128

N = 256

N = 512

(c) Last Layer Kernel Consistency

10 2

3×10 1

4×10 1

6×10 1

2×10 2

N

10

1

1 - CKA(

4

N

,

4

512

)

Final vs

N

N

1

Init vs Final

(d) Kernel Alignment vs

N

N

=64

=1

=2

=3

=4

N

=128

N

=256

N

=512

In this section, we show additional figures illustrating convergence of network quantities across widths that we did not have space for in the main text.

B.1 Vision

A simple setting in which convergence properties are particularly clear and simple to study is for a ResNet learning CIFAR-10 and going over multiple passes of the dataset. In Figure 10 we plot a 20-epoch pass over CIFAR 10, and study the generalization error, initial and final preactivations in the last layer, and final layer kernels across widths. The training error begins to exhibit pathologies after sufficiently many epochs, related to the discussion in section 3.2.

Next in figure 11, we show a higher-resolution plot of the kernel Gram matrices across widths and across layers for the CIFAR-5M ResNet after a pass through the data. The larger resolution allows one to see that even the fine-grained details in the structure of the Gram matrix are consistent across widths.

B.2 Language

Next, in Figure 12, we create an analog of the language column of Figure 3, this time for µ P transformers of the same architecture and dataset but now optimized with vanilla SGD. The fact that wider transformers perform better still holds, and one can clearly see narrower networks approaching wider ones in their output logit values.

10 0

10 1

10 2

% of Wikitext-103 seen

3.5

4.0

4.5

5.0

5.5

6.0

6.5

7.0

7.5

Training Errors

Single network

Ensembled

d=64 sing

d=128 sing

d=256 sing

d=512 sing

d=1024 sing

d=2048 sing

d=4096 sing

(a) Training Loss

0.0

0.2

0.4

0.6

0.8

1.0

Tokens seen

1e8

6

7

8

9

10

11

12

13

14

15

True logit value on test point

N=64

N=128

N=256

N=512

N=2048

N=4096

(b) Correct logit value

0.0

0.2

0.4

0.6

0.8

1.0

Tokens seen

1e8

0.05

0.10

0.15

0.20

0.25

0.30

True logit RMSE from

N

=4096

N=64

N=128

N=256

N=512

N=2048

(c) RMSE from widest

C Defining µ P and SP (Standard Parameterization)

There are several detailed discussions about µ P vs SP scaling [12, 15, 9, 10, 14]. The aim of this section is to simply give an accessible and conceptual overview of their distinction, as well as a motivation for µ P from the perspective of keeping features moving in time even at infinite width. .

There are several equivalent ways of parameterizing neural networks that give rise to the same dynamical effects, whether in µ -parameterization or standard parameterization. We give the definitions in the case of a single-output feed-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 a detailed discussion see [12] and also [14].

C.1 SP

We assume all hidden layers have equal width N . Let the input space have dimension D . Let µ be the index of the training point in the dataset. At each layer ℓ , the pre-activation h ℓ +1 µ in layer ℓ +1 is given by

where ϕ is an element-wise non-linearity, often taken to be the ReLU function. Here the N -1 / 2 out front allows h ℓ +1 µ to be Θ(1) at initialization as N →∞ by the law of large numbers. The output of the network f µ is then given by:

Here again the N -1 / 2 scaling again yields that f µ will be Θ(1) as N →∞ . In SP, α is taken to be 1 , but we will keep it explicit as it plays an important role in distinguishing the parameterizations. It is the laziness parameter identified in [31]. The change in the function is given by

Here K µν = ∇ θ f µ · ∇ θ f ν is the NTK gram matrix. y ν is the true label. ℓ is the loss function (e.g. MSE or crossentropy) and ℓ ′ is its derivative with respect to the first argument. The NTK is easily seen to be order α 2 and ℓ ′ is order 1 at small α . In order to have the change in the function be Θ(1) we set η = η 0 /α 2 .

Using the chain rule, one can directly see that the pre-activations evolve as [10, 3, 32]

Thus, at large N and α = 1 the pre-activations of this network evolve as Θ( N -1 / 2 ) . Consequently, at infinite width the feature do not evolve and infinitely wide networks in standard parameterization become kernel machines with the static and initialization-independent infinite-width NTK.

In many machine learning libraries, the factors of 1 / √ N are not explicitly placed in front of each multiplication with the weight matrices. Rather, the weight matrices themselves are drawn from a distribution W ℓ ∼ N (0 , 1 N 1 ) . Although this gives identical forward pass, this changes the N scaling of the gradients in the backward pass. As long as the learning rate is appropriately rescaled to account for this, the dynamics are equivalent to the SP parameterization discussed above.

C.2 µ P

One of the simplest ways to define the µ -parameterization is to take α = 1 / √ N . This implies that we simply replace the final layer of the network by:

As the prior analysis shows, in order to have df µ /dt be Θ(1) at initialization, we take η = Nη 0 , so the learning rate in this definition scales extensively with N . In this setting, we now have that

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

D Importance of Parameterization

1

2

3

4

5

6

Epochs

1.2

1.4

1.6

1.8

2.0

Test Loss

N = 128

N = 128 NTK param.

N = 256

N = 512

N = 1024

0

1

2

3

4

5

6

Epochs

0.0000

0.0005

0.0010

0.0015

0.0020

Kernel Alignment Increase

N = 128

N = 256

N = 512

N = 1024

In this section, we aim to illustrate that many of the claims of width consistency at practical widths are in fact contingent on using the mean field/ µ -parameterization. In Figure 13 we compare the early dynamics of network training on CIFAR-10 between networks in the NTK parameterization (dashed lines) and in mean field parameterization (solid). The models are initialized in such a way so that their N = 128 dynamics perfectly coincide. We plot both test error and the kernel alignment between the NTK and the target function. While both parameterizations will eventually converge to an infinite width limit, we note the following two differences in width scaling behavior:

Wider networks tend to train more slowly in NTK parameterization and perform worse as the width grows. This is because the rate of feature learning is not held constant across widths and decreases as width grows. Wider networks in µ P train slightly faster.

Mean field networks approach their limit more quickly as width increases in terms of both test loss and kernel dynamics.

The general theoretical principle behind these observations is that mean field parameterization generates feature updates which are scale independent, thus eliminating an unnecessary source of finite size approximation error in the dynamics [39].

E Further Studies of µ P versus SP

E.1 MLPs

In figure 14, we show a 3-layer MLP learning a quadratic polynomial. In subfigure a) use a batch size of 10 and a learning rate going as η = 5 N/ (1 + α 2 0 ) with α 0 = 1 . The output layer is scaled as α 0 / √ N , putting us in the rich regime. The learning rate has been picked to be nearly as large as possible at this batch size in order to maximize the large loss curve fluctuations yielded by large learning rate effects. In subfigure c) we do not rescale the output layer, and have a width-indepdent learning rate going as η = 50 / (1 + α 2 0 ) with α 0 = 1 . This puts is in the large-learning rate regime for a standard parameterized network. See Appendix A for more details.

We plot the learning curves across widths and find striking agreement, even at the fine-grained level of fluctuations due to small batch size and large learning rate effects. Although this is not exactly the full-batch edge-of-stability effect reported in [7], the large oscillations may be similar to a small batch size analog. We plot the absolute difference from the widest network in subfigure b) to highlight the strong agreement across widths.

In subfigure c), we have the same network but in standard parameterization. The narrower networks now learn features more quickly, leading to inconsistent dynamics across widths.

0

50

100

150

200

250

300

350

400

Steps of SGD

10

2

10

1

10 0

Train Loss

<

>

N=100

N=223

N=499

N=1118

N=2499

0

50

100

150

200

250

300

350

400

Steps of SGD

10

5

10

4

10

3

10

2

10

1

Absolute difference of train loss from widest

N=100

N=223

N=499

N=1118

0

50

100

150

200

250

300

350

400

Steps of SGD

10

2

10

1

10 0

Train Loss

<

>

N=100

N=223

N=499

N=1118

N=2499

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 to that figure, we see significantly different dynamics and batch variation across widths. In Figure 15 a) we plot the early time behavior of a CIFAR-5m task at large learning rate. The large learning rate effects cause the loss to substantially oscillate, but the oscillations across widths are inconsistent by contrast to 6a b) . Further, at late times in Figure 15, the sharp spikes in the loss function due to large learning rate effects become substantially different across widths. Indeed, in SP some widths converge for a given learning rate while others do not. This trend has already been well-studied in [15]. We again stress that our observation is that not only are the final losses similar across widths in µ P (as observed in [15]), but that the individual batch and large learning rate fluctuations agree across widths at early times in µ P as well.

20

25

30

35

40

45

50

55

60

Train Steps

2.2

2.3

2.4

2.5

2.6

Train Loss

64

128

256

0

10

20

30

40

50

60

Train Steps

0

20

40

60

80

100

Train Loss

64

128

256

E.3 Language

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

10 0

10 1

10 2

% of Wikitext-103 seen

3.5

4.0

4.5

5.0

5.5

6.0

6.5

7.0

7.5

Training Errors

Single network

Ensembled

d=64 sing

d=128 sing

d=256 sing

d=512 sing

d=1024 sing

d=2048 sing

d=4096 sing

(a) Training Loss

0.0

0.2

0.4

0.6

0.8

1.0

Tokens seen

1e8

3

4

5

6

7

8

9

10

11

12

True logit value on test point

N=64

N=128

N=256

N=512

N=2048

N=4096

(b) Correct logit value

0.0

0.2

0.4

0.6

0.8

1.0

Tokens seen

1e8

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

True logit RMSE from

N

=4096

N=64

N=128

N=256

N=512

N=2048

(c) RMSE from widest

F Bias-Variance Decompositions over Initializations

In this section we explicitly define what we mean by initilization bias and initialization variance. Following [30], we consider the trained neural network function f θ ∗ ( x ) to depend on the training set D (including inputs x , outputs y , and possible label noise ϵ ) as well as the initial parameters θ 0 . Here, θ 0 are the initial parameters and θ ∗ 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 ( x µ , y µ ) P µ =1 . By contrast, our paper focuses on the variance due to the initial parameters θ 0 from which training begins. In the overparameterized regime with more parameters than data-points and relatively little label noise, this has been shown to be the dominant source of variance for neural networks [27, 29, 28].

Consequently, our definition of the bias of a neural network is given by

The bias can be approximated by averaging a sufficiently large ensemble of neural networks over initialization seeds. Each network in the ensemble is trained on the same dataset in the same batch order using the same optimizer.

The variance of the neural network predictor is given by:

The variance of E ensembles of a given network with independent initialization seeds θ e is given by

As E →∞ , the empirical average of E network ensembles approaches the bias, so the variance of the ensembled networks goes to zero. As long as the errors are uncorrelated, the variance the the network ensembles decreases as 1 /E . In practice, we see that comparable decay rates are achieved.

G Task-Dependent Scaling Laws in Width

For the settings where we observed larger deviations in the dynamics for models of varying widths, we examined scaling of the training losses with respect to width N after a significant amount of training (2 epochs for ImageNet and 1 epoch for Wikitext 103). We fit power laws of the form Loss = CN -β + D where C, β, 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 2 above .99.

The existence of such power law behavior across widths provides further evidence of the networks approaching a well-defined inifinite width limit. One could imagine that the differences between successive networks might get smaller (as in Figure 4) but that there is no well-defined limit as N →∞ , similar to the terms in a harmonic sum. The fact that the power law fit has exponent much larger than 0 and does not display logarithmic dependence on N provides empirical evidence that we expect convergene as N →∞ .

Additionally, given that the observed power laws N -β are task-dependent and significantly different from N -1 suggests that models at late time are not well described by perturbation theory around the infinite width mean field limit, which predicts a universal exponent of β = 1 [54, 39]. This motivates novel theoretical descriptions of finite width mean field

10 2

10 3

N

10 0

2×10 1

3×10 1

4×10 1

6×10 1

Loss - D

Loss =

CN

+

D

Data, Wikitext

6.95 ×

N

0.42

Data, Imagenet

11.2 ×

N

0.63

learning networks at late time which can capture task-dependent exponents.

H Overview of Finite Width Corrections to Feature Learning Networks

In this section we review some basic ideas from the mean field theory of feature learning neural networks. We first describe the predictions that mean field theory makes about infinite width networks before describing finite size corrections to the dynamics of learning. To eliminate unnecessary complexity, we will focus on MLP layers, but these arguments can be easily extended to CNN and self-attention layers as well. We start by defining a MLP in a parameterization equivalent to µ P

We will consider these networks trained from a random Gaussian initialization of the weights so that θ = Vec { w L , ..., W 0 } follows θ ∼ N (0 , I ) at initialization. This network is then trained with some gradient based optimizer, leading to dynamical predictions f µ ( t ) and dynamical preactivations h ℓ µ ( t ) . Because of the random initialization of weights, the outputs of the network and the precise preactivations are random variables. However, at infinite width N →∞ , a dramatic simplification of the dynamics occurs.

H.1 The Infinite Width/Mean Field Limit

The predictions f µ ( t ) and internal representations of infinite width limit of neural networks admit a description in terms of non-random initialization-independent dynamical feature kernels Φ ℓ µν ( t, s ) and gradient kernels G ℓ µν ( t, s ) defined as

where µ, ν index data points and t, s index training time and g ℓ µ ( t ) = N ∂f µ ∂ h ℓ are back-propagated gradient signals [12, 14]. Further, all preactivation vectors h ℓ µ ( t ) ∈ R N have entries that become iid draws from a (potentially non-Gaussian) single site density p ( h ) , which converges as

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 Φ ℓ µν ( t, s ) = ∫ p ( h ℓ µ ( t ) , h ℓ ν ( s )) ϕ ( h ℓ µ ( t )) ϕ ( h ℓ ν ( s )) dh ℓ µ ( t ) dh ℓ ν ( s ) .

H.2 Finite Width Effects

At finite width, the internal kernels { Φ ℓ µν ( t, s ) , G ℓ µν ( t, s ) } and predictions f µ ( t ) of the model become initialization and width-dependent and deviate from their mean field dynamics. For Gaussian random initialization of the weights of the network, the predictions and kernels fluctuate (from init to init) with variance that scales asymptotically like Θ(1 /N ) for width N (or 1 /d model for transformer) [39]. Further, the ensemble averaged values for the predictions ⟨ f µ ( t ) ⟩ and kernels 〈 Φ ℓ µν ( t, s ) 〉 differ asymptotically from their infinite width values by Θ( N -1 ) . Both of these two leading order effects can influence the expected (train or test) loss of the model. At fixed width and late training time, finite size effects beyond leading order can accumulate and become relevant, however theory predicts that any observable average at width N admits an asymptotic series in powers of N -1 [39].

H.2.1 Trainability at Finite Size

The Θ( N -1 ) correction to feature and gradient kernels can lead to non-trivial corrections to the loss dynamics. Working in continuous time, we can define the neural tangent kernel (NTK) as K µν ( t ) = ∑ ℓ G ℓ +1 µν ( t, t )Φ ℓ µν ( t, t ) , where base cases are Φ 0 µν ( t, s ) = 1 D x µ · x ν and G L +1 µν ( t, s ) = 1 . Following the approximation to online dynamics with MSE loss in Section 4, we consider a gradient flow on the average dynamical NTK

where T is the time-ordering operator. We now consider the leading correction to the average NTK around infinite width ⟨ K ( t ) ⟩ = K ∞ ( t ) + 1 N K 1 ( t ) + Θ( N -2 ) . With this correction, we see that the dynamics of errors ∆

The fact that the 1 N K 1 correction is integrated over time and placed in the matrix exponential indicates that small corrections to NTK dynamics can lead to large dynamical amplification of logit corrections. This fact was pointed out in another work [39] which tried to motivate a study of perturbation theory in logarithms of the transition matrix log T ( t ) defined as

The solution to this can be used to construct the errors at a later time ∆ ( t ) = exp ( R ( t )) y .

I The empirical sufficiency of ensembling only a few times

Our experiments rely on ensembles of small numbers of neural networks to analyze the bias component of the loss as it varies across width. One can show the marginal value of adding a network to the ensemble decreases with the ensemble size. Figure 18 illustrates this in the setting of ResNets trained on ImageNet. In Figure 18(a), we show that the reduction in variance over initializations due to ensembling rapidly plateaus as soon as the ensemble size reaches E = 3 . In Figure 18(b), we show that the loss curves as function of width are very similar for ensemble sizes above 3. Indeed, the green, orange, and red curves - corresponding to E ∈ { 3 , 4 , 5 } -are nearly identical.

Lastly, Figure 18(c) confirms that the initialization variance plays a negligible role in the scaled RMSE distances between true logits across width, even for single draws of networks trained on a small number of examples. The dashed lines in this Figure correspond to scaled RMSE distances between networks of the same widths as their solid-line companions, albeit with weights initialized to zero. The ordering and scale of the curves are comparable, and in fact nearly identical for the curves corresponding to N ∈ { 128 , 256 } compared to N = 512 .

1

2

3

4

5

Ensemble size

2.2

2.4

2.6

2.8

3.0

3.2

3.4

Test loss

N

=32

N

=64

N

=128

N

=256

N

=512

(a) Test loss vs.

#

ensembles

32 64

128

256

512

Base width (

N

)

2.2

2.4

2.6

2.8

3.0

3.2

3.4

Test loss

E

=1

E

=2

E

=3

E

=4

E

=5

(b) Test loss vs. width

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Epochs

0.0

0.2

0.4

0.6

0.8

1.0

True logit RMSE from

N

=512

zeroed

32

zeroed

64

zeroed

128

zeroed

256

(c) Test loss for subtracted networks

J Offline Training

Figure 19 depicts the loss curve for a ConvNeXt-T (tiny) model trained on ImageNet in the typical, offline setting - where data is encountered repeatedly across many epochs. As the networks overfit the training data - in Figure 19, beyond 40,000 training steps or five epochs - the loss curves diverge dramatically for different-width networks. Width consistency subsequently erodes.

100

1000

10000

Train Steps

0.5

1.0

2.0

4.0

8.0

Train Loss

48

96

192

384

K Use of Compute

For most experiments, we used Nvidia A100 SXM4 40GB and 80 GB GPUs on an academic cluster.

For the Wikitext-103 tasks, each width included 4 ensembles loaded onto an A100 GPU that ran for a range between 1 to 3 days. For each sweep over widths this corresponds to about 8 A100-days. Accounting for sweeps over different sequence lengths, optimizers, and parameterizations, this corresponds to about 50 A100-days.

All MLP tasks, including the calculation of empirical NTKs and their spectral properties were done in 15-30 minute Colab sessions using the basic GPUs provided.

The CIFAR-10 ResNet experiments in Figure 10 were done using a total of less than 1 A100-day of compute across all widths and ensembles.

The ImageNet ResNet experiments vectorize training over between one to four same-width neural networks on one A100 GPU. Each experiment training a collection of networks for 30 epochs takes between one to three A100-days. Overall, these experiments expended roughly 30 A100-days.

For the CIFAR-5m experiments in Figure 2 and 3, across all widths, it required a few hours of A100 GPU. For Figure 7 and 8, as these were ensembled across multiple runs, these required close to 1-2 days of A100-GPUs. Figure 6a was just run for a few 100 steps of the training, so didn't use much compute power.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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**\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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)\n\nQUESTIONS:\nPlease 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)\n\nLIMITATIONS:\nThere isn't much discussion of limitations within the submission.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThere 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nHow 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nhe 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.\n\nQUESTIONS:\nIn 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\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nMost 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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nGiven 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.\n\nWEAKNESSES:\nNote 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\n\nQUESTIONS:\n(Please see the comments above.)\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-26"}