Title: Infinite Limits of Multi-head Transformer Dynamics
Venue: NeurIPS 2024
Accepted: yes
Keywords: infinite limits

transformers

mean field theory

learning dynamics
Abstract: In this work we analyze various scaling limits of the training dynamics of transformer models in the feature learning regime. We identify the set of parameterizations which admit well defined infinite width and depth limits that allow the attention layers to update throughout training, a relevant notion of feature learning in these models. We then use tools from dynamical mean field theory (DMFT) to analyze various infinite limits (infinite heads, infinite key/query dimension, and infinite depth) which have different statistical descriptions depending on which infinite limit is taken and how attention layers are scaled. We provide numerical evidence of convergence to the limits and show they maintain the correct scale of updates for both SGD and Adam.

Infinite Limits of Multi-head Transformer Dynamics

Abstract

In this work, we analyze various scaling limits of the training dynamics of transformer models in the feature learning regime. We identify the set of parameterizations that admit well-defined infinite width and depth limits, allowing the attention layers to update throughout training-a relevant notion of feature learning in these models. We then use tools from dynamical mean field theory (DMFT) to analyze various infinite limits (infinite key/query dimension, infinite heads, and infinite depth) which have different statistical descriptions depending on which infinite limit is taken and how attention layers are scaled. We provide numerical evidence of convergence to the limits and discuss how the parameterization qualitatively influences learned features.

1 Introduction

Increasing the scale of transformer models has continued to improve performance of deep learning systems across many settings including computer vision [1, 2, 3, 4] and language modeling [5, 6, 7, 8, 9]. However, understanding the optimization stability and limiting behavior of these models under increases in model scale remains a core challenge.

One approach to scaling up systems in a stable and predictable way is to identify parameterizations of neural networks that give approximately scale-independent feature updates during training [10, 11, 12]. The mean field parameterization, commonly referred to as µ P, is a well-known example that satisfies this property [13, 14, 15]. When such parameterizations are adopted, the learned internal representations in hidden layers of the network are very similar across model scales [16, 17], but performance tends to improve with model scale [10, 11, 12]. Further, theoretical results about their limits can often be obtained using Tensor Programs [14] or dynamical mean field theory (DMFT) techniques [15, 17].

In this work, we develop a theoretical treatment of randomly initialized transformers. We study various scaling limits of the training dynamics of these models including the infinite key/query dimension limit, the infinite head limit, and the infinite depth limit. Concretely, our contributions are the following:

We derive a DMFT for feature learning in randomly initialized transformers with key/query dimension N , attention head count H and depth L . From the derived DMFT action, we identify large N , large H and large L limits of the training dynamics.

We analytically show that the large key-query N → ∞ limit requires the µ P scaling of key/query inner product with 1 /N , even if key/queries are reparameterized to decrease the size of their updates from gradient descent.

From the limiting equations, we show that this N →∞ limit causes multi-head self attention trained with stochastic gradient descent (SGD) to effectively collapse to single-head self attention since all heads follow identical dynamics.

To overcome this limitation, we analyze the infinite head H → ∞ limit while fixing N . We show there is a limiting distribution of attention variables across heads at each layer throughout training. Despite N being finite, the infinite-head H → ∞ limit leads to concentration of the network's output logits and learned residual stream feature kernels, giving deterministic training dynamics.

Finally, we examine large depth limits of transformers with residual branch scaling. We illustrate and discuss the tension between parameterizing a model so that it has a non-trivial kernel at initialization while maintaining feature learning within the multi-head self attention (MHSA) and multi-layer perceptron (MLP) blocks.

1.1 Related Works

Hron et al. [18] studied the Neural Network Gaussian Process limit of multi-head self attention in the infinite-head H → ∞ limit. They showed that, at initialization, there is a limiting distribution over attention matrices and that the outputs of the multi-head attention block follow a Gaussian process, establishing a connection to kernel methods. Dinan et al. [19] develop a similar theory of transformers at initialization and compute the Neural Tangent Kernel associated with this architecture as the dimensions per head N → ∞ using a 1 √ N scaling of the key-query inner product within each attention layer. One of our key theoretical results is showing that this picture of a distribution over learned attention heads persists throughout training in the feature-learning regime as H → ∞ (though the distribution of residual stream variables generally becomes non-Gaussian).

Several works have analyzed the signal propagation properties of transformers at initialization at large key/query dimension N and large depth L [20, 21, 22, 23] including providing modifications to the standard transformer architecture [22, 24]. In this work, we pursue large depth limits of transformers by scaling the residual branch as L -α L with α L ∈ [ 1 2 , 1] , which has been shown to converge to a limit not only at initialization [25, 26, 27], but also throughout training in the feature learning regime [11, 12, 27]. However, we argue that in transformers that α L = 1 is preferable as it enables the attention layers to update non-negligibly as L →∞ .

Yang et al. [10] introduced the µ P scaling for attention layers which multiplies the key/query inner product with 1 N rather than the more commonly used 1 √ N [5]. They show empirically that this change improves stability of training and transfer of optimal hyperparameters across different values of N . Vyas et al. [16] empirically found that such µ P transformers learn attention matrices that become approximately consistent across different heads and model sizes, suggesting that models parameterized in µ P learn similar representations across scales.

In addition to work on infinite width and depth limits of deep networks, there is also a non-asymptotic approach to optimizer design and scaling based on controlling the norm of weight updates [28]. This approach coincides with µ P width-scaling when the spectral norm of the weights is used as the measure of distance [29], and can achieve hyperparameter transfer for a wide array of optimizers and initialization schemes [30 ? ].

2 Parameterizations with Feature Learning Limits

We consider a transformer architecture with L layers, H heads per layer, and N dimensional keys/ queries per head. Transformers are often defined in terms of d model = H d head = H N which can be increased by scaling the number of heads or the dimension of each head, where N is often written d head . Our goal is to determine the set of parameterizations that allow the attention layers to undergo non-trivial feature learning in the various N, H , L → ∞ limits. The analysis of these limits is performed with batch size and number of training steps t fixed while the other architectural parameters are taken to infinity.

2.1 Model Scalings

The network's output is computed by a depth L recursion through hidden layers ℓ ∈ [ L ] starting with the first layer h 1 s ( x ) = 1 √ D W 0 x s ∈ R N H where x s ∈ R D is the input at spatial/token position s . Preactivations in subsequent layers h ℓ are determined by a forward pass through the residual stream which contains an attention layer and a MLP layer

The constants γ 0 and β 0 control the rate of feature learning and the effective depth respectively 1 . We will consider α L ∈ [ 1 2 , 1] . 2 The multi-head self attention layer (MHSA) with pre-layer-norm 3 is where σ ℓ h ∈ R S×S are the attention matrices passed through a matrix-valued nonlinearity σ ( A ℓ h ) 4 . For a sequence of length S , the pre-attention matrix A ℓ h ∈ R S×S is defined as

The exponent α A will alter the scale of the pre-attention variables A ℓ h at initialization. The input matrices have shape W ℓ V h , W ℓ K h , W ℓ Q h ∈ R N × N H , while the output matrices have shape W ℓ O h ∈ R N H× N . All of the weights W ℓ O h ′ , W ℓ Q h , W ℓ K h are initialized with Θ(1) entries while W ℓ K h , W ℓ Q h have entries of size Θ( N 1 -α A ) which ensures that all key and query k , q vectors are Θ(1) at initialization. The pre-attention variables A ℓ h ∈ R S×S at each head h are determined by key k ℓ hs and query q ℓ hs ′ inner products. The MLP layer consists of two linear layers with an element-wise nonlinearity ϕ applied in between, where W ℓ, 2 , W ℓ, 1 ∈ R N H× N H are initialized with Θ(1) entries:

µ P scaling [13, 31, 14, 15] downscales the readout of the last layer compared to standard and NTK parameterization [32]. Thus, we define the output of the model as

5 where h L s ∈ R N H are the final layer preactivations at spatial/token position s ∈ [ S ] . The parameter γ 0 is an additional scalar that controls the rate of change of the internal features of the network relative to the network output [33].

2.2 Learning Rate Scalings

In order to approximately preserve the size of internal feature updates, we must scale the learning rate η appropriately with ( N, H , L ) . However, this scaling depends on the optimizer. In Table 2, we provide the appropriate scaling of learning rates for SGD and Adam for any α L ∈ [ 1 2 , 1] and α A ∈ [ 1 2 , 1] . In what follows, we focus on the SGD scaling and theoretically analyze the N →∞ , H → ∞ , and L → ∞ limits of the training dynamics. We also provide in Table 2 details about what prefactor the first layer should be multiplied by and the initial weights divided by to ensure convergence to the L →∞ limit. Example FLAX implementations of these parameterizations for vision and language modeling transformers are provided in Appendix B.

(a) Scaled Residual Stream

(b) MHSA Block

Ah=

1

kh·qh

NαA

HeadsH

MHSA

MLP

×

LaL

βo

MHSA

×

LaL

βo

MLP

LaL

βo

+

·..N

Key

Queries

Attention

Values

Output

H

Our analysis assumes that at each step t of SGD or Adam a mini-batch B t of size Θ(1) is used to estimate the loss gradient. We assume that the minibatches are fixed. Further, the number of total training steps is assumed to not be scaled jointly with the model size. Therefore the analysis provided here can cover both online SGD for a fixed number of steps or full batch GD (repeating data) with a Θ(1) sized dataset.

3 Infinite Limits of Learning Dynamics

In this section, we first analyze the infinite dimension-per-head N →∞ limit of training. We find that for this limit, the µ P rule of α A = 1 is necessary and show that all heads collapse to the same dynamics. To counteract this effect, we next analyze the infinite head H → ∞ limit of the training dynamics at fixed N,L , where we find a limiting distribution over attention heads. We will conclude by analyzing the statistical descriptions of various infinite depth L →∞ limits. 6 .

3.1 Mean Field Theory Treatment of the Learning Dynamics

To obtain the exact infinite limits of interest when scaling dimension-per-head N , the number of heads H , or the depth L to infinty, we work with a tool from statistical physics known as dynamical mean field theory (DMFT). Classically, this method has been used to analyze high dimensional disordered systems such as spin glasses, random recurrent neural networks, or learning algorithms with high dimensional random data [34, 35, 36, 37, 38, 39]. Following [15, 11], we use this method to reason about the limiting dynamics of randomly initialized neural networks by tracking a set of deterministic correlation functions (feature and gradient kernels) as well as additional linear-response functions (see Appendix D). The core conceptual idea of this method is that in the infinite limit and throughout training, all neurons remain statistically independent and only interact through collective variables (feature kernels, neural network outputs, etc). Further the collective variables can be computed as averages over distribution of neurons in each hidden layer or along the residual stream. This DMFT description can be computed using a path integral method that tracks the moment generating function of the preactivations or with a dynamical cavity method (see Appendix D).

3.2 Scaling Dimension-Per-Head N

One way of obtaining a well-defined infinite parameter limit of transformers is to take the N →∞ limit, where N is the dimension of each head. A priori, it is unclear if there are multiple ways of scaling the key/query inner product. Concretely, it is unknown what values for the exponent α A are admissible for the pre-attention A = 1 N α A k · q . The keys and queries are uncorrelated at initialization which motivated the original choice of α A = 1 2 [5, 18]. Yang et al. [10] assume the entries of the key and query vectors move by Θ(1) , implying α A = 1 is necessary since the

10

2

10

1

10 0

Learning Rate

0.8

1.0

1.2

1.4

1.6

Test Loss

=1

= 1

2

N

=4

N

=8

N

=16

N

=32

N

=64

(a) Hyperparameter Transfer for Various α A

10 1

10 2

10 3

N

10

6

10

4

10

2

10 0

10 2

Var

N

2

=1

= 1

2

(b) Attention variance across heads

update to k is correlated to q and vice versa. However, it is possible to obtain Θ(1) updates to the attention variable for alternative values of α A if we choose the change to key (also query) entries after gradient descent to be δk i ∼ Θ( N -1+ α A ) . We show that this scaling can approximately preserve optimal hyperparameters across N in Figure 2 (a) and give similar dynamics under SGD Appendix C. However, as we show in Appendix E.1.2, any well defined N →∞ limit of SGD requires α A = 1 . The reason is not that keys and queries become correlated, but rather that the scale of the backward pass must be controlled to ensure the dynamics remain stable (non-divergent) under SGD training. After performing two or more gradient descent steps, we demonstrate that the backpropagation signals will diverge as N →∞ unless initial key and query weight matrices are downscaled to have variance of order Θ N (1) . In Appendix E, we provide a DMFT analysis of the N →∞ limit of the transformer training dynamics. We summarize the result of that analysis informally below.

Result 1 (Infinite Dimension-Per-Head N ) ( Informal ) A stable feature learning N →∞ limit of transformer SGD training requires taking α A = 1 ( µ P scaling), even if key/query updates are allowed to be rescaled to account for their correlation. The limiting dynamics of training are governed by the residual stream kernel H ℓ ss ′ ( x , x ′ , t, t ′ ) = 1 N H h ℓ s ( x , t ) · h ℓ s ′ ( x ′ , t ′ ) and a collection of inner product kernels in each head h that concentrate as N →∞

alongside residual-stream gradient kernels and response functions in the sense of [15, 11]. The NN output logits f ( x , t ) evolve deterministically according to the above kernels as well as kernels for the gradient vectors g ℓ ≡ γ 0 N H ∂f ∂ h ℓ which appear in the backward pass. These variables become identical across heads such that for any h , h ′ ∈ [ H ] , A ℓ hss ′ ( x , t ) = A ℓ h ′ ss ′ ( x , t ) . All preactivations on the residual stream and key/query/value variables within a MHSA block are statistically independent across neurons and can be described by a single scalar stochastic process

where ˜ u ℓ , u ℓ , u ℓ K h are Gaussian processes with covariances Φ ℓ, 1 , V ℓσ , H ℓ respectively. Analogous equations hold for the queries and values. In the limit, the kernels H ℓ ss ′ ( x , x ′ , t, t ′ ) = 〈 h ℓ s ( x , t ) h ℓ s ′ ( x ′ , t ′ ) 〉 , A ℓ hss ′ ( x , t ) = 〈 k ℓ hs ( x , t ) q ℓ hs ′ ( x , t ) 〉 , etc. are computed as averages ⟨·⟩ over

these random variables. The deterministic kernels C ℓ , ˜ C ℓ can also be expressed in terms of single site averages of residual variables and head averages of MHSA variables. The kernels C ℓ , ˜ C ℓ , C k ℓ depend on the precise mini-batches of data B t presented at each step t which we assume are known.

We derive this result using a Martin-Siggia-Rose path integral formalism [40] for DMFT in Appendix E. Full DMFT equations can be found in Appendix E.2. Following prior works on DMFT for infinite width feature learning, the largeN limit can be straightforwardly obtained from a saddle point of the DMFT action [15, 11, 41, 17].

Collapse of Attention Heads As N →∞ , multi-head self-attention will effectively compute the same outputs as a single-head self-attention block. We theoretically derive this effect in Appendix E.2.1 and demonstrate it empirically in Figure 2 (b). This experiment shows that in α A = 1 ( µ P) transformers trained for 2500 steps on CIFAR-5M [42], the variance of attention matrices across heads decreases with N . However, we note that if the scaling exponent is chosen instead as α A = 1 2 there is non-decreasing diversity of attention variables across heads. This result is consistent with recent empirical findings that attention variables in µ P transformers converge to the same limiting quantities at large N with H fixed for different initializations and also across model sizes [16]. This aspect of transformers in the largeN limit is potentially undesirable as some tasks could require learning multiple attention mechanisms. Furthermore, this suggests scaling the model in this limit could increase computational cost without improving performance. To circumvent this, we explore if there exist well defined limits at finite N with a diversity of attention variables across heads.

3.3 Scaling Number of Heads

In this section, we take H → ∞ with the inner dimension N fixed. Rather than having all kernels concentrate, the kernel of each head of the MHSA block follows a statistically independent stochastic process. This picture was shown to hold at initialization by Hron et al. [18]. Here, using a DMFT analysis, we show that it continues to hold throughout training in the feature learning regime.

Result 2 (Infinite Head Limit) ( Informal ) The H → ∞ limit of SGD training dynamics in a randomly initialized transformer at any key/query dimension N , scaling exponents α A , α L , and any depth L is governed by head-averaged kernels for pairs of input data x , x ′ at training times t, t ′ and spatial/token positions s , s ′ such as

which converge to deterministic values as H → ∞ . The attention variables { k ℓ h ( x , t ) , q ℓ h ( x , t ) , v ℓ h ( x , t ) , A ℓ h ( x , t ) } within each head become statistically independent across heads and decouple in their dynamics (but not across dimensions within a head). Each residual stream neuron becomes independent and obeys a single site stochastic process analogous to Result 1, but with different kernels.

We derive this and the full DMFT in Appendix E.3, showing that the joint distribution of headaveraged dynamical quantities satisfies a large deviation principle and the limit can be derived as a saddle point of a DMFT action.

To gain intuition for this result, we first examine variables H ℓ and A ℓ h at initialization. In Figure 3, we plot the convergence of a N = 4 , L = 8 vision transformer's residual stream kernel H ℓ to its H → ∞ limit at rate O ( H -1 ) in square error, consistent with perturbative analysis near the limit [17]. Next, we plot the distribution (over heads) of A h at a fixed pair of spatial/token positions for a fixed sample. This is a non-Gaussian random variable for finite N , but as N →∞ the distribution of A will approach a Gaussian with variance Θ( N 1 -2 α A ) .

We then investigate training dynamics as we approach the H → ∞ limit. In Figure 4 (a) we show the test loss on CIFAR-5M as a function of the number of training iterations. The performance tends improve as H increases and the model approaches its limit. In Figure 4 (b) we show that all of the models are converging in function space by measuring the squared error between finite H head models and a proxy for the infinite H model. Since the H → ∞ limit is essentially uncomputable, we approximate it as the ensemble averaged predictor of the widest possible models, a technique used in prior works [16, 11]. We again see that at early time, the logits of H head models converge to the limit proxy at a rate O ( H -1 ) , but after continued training the convergence rate weakens. This effect

10 1

10 2

10 3

10

5

10

4

10

3

10

2

|

H

H

|

2

1

=1

=2

=3

=4

=5

=6

=7

=8

(a) Initial

H

ℓ

Kernels,

N

= 4

2

1

0

1

2

10

2

10

1

10 0

Density

=32

=64

=128

=256

(b)

A

h

Distribution

N

= 1

2

1

0

1

2

10

2

10

1

10 0

Density

=32

=64

=128

=256

(c)

A

h

Distribution

N

= 16

0

500

1000

1500

2000

2500

Steps

1.0

1.2

1.4

1.6

1.8

2.0

Test Loss

=16

=32

=64

=128

=256

=512

=1024

=2048

(a) Training Dynamics Varying

H

10 1

10 2

10 3

10

4

10

3

10

2

10

1

10 0

|

f

f

|

2

1

steps=25

steps=100

steps=250

steps=500

steps=1000

steps=2000

(b) Convergence to

H → ∞

limit

has been observed in µ P networks in many settings [16] and a theoretical model of this was provided in recent work which argues it arises from low-rank effects in the finite H kernels [39].

3.4 Infinite Depth Limits

We next describe the infinite depth limits which depend on the choice of α L . Below we informally describe the main finding which again uses a DMFT formalism and is based on analyses in recent works on infinite depth networks from Bordelon et al. [11] and Yang et al. [12].

Result 3 (Infinite Depth Limit) ( Informal ) The training dynamics for H , L → ∞ with L -α L branch scaling with α L ∈ [ 1 2 , 1] is described by a differential equation for residual variables h s ( τ, t ) in layer time τ = lim L →∞ ℓ L for the residual stream

where the Brownian motion term du s ( τ, x , t ) survives in the limit only if α L = 1 2 and has covariance ⟨ du s ( τ, x , t ) du s ′ ( τ ′ , x ′ , t ′ ) ⟩ = δ ( τ -τ ′ ) dτdτ ′ [Φ ss ′ ( τ, x , x ′ , t, t ′ ) + V σ ss ′ ( τ, x , x ′ , t, t ′ )] (11)

and the deterministic kernel C ss ′ ( τ, x , x ′ , t, t ′ ) can be expressed in terms of head-averaged kernels and response functions. The weights inside each hidden MHSA layer or each MLP layer are frozen in the L →∞ limit unless α L = 1 . All response functions are suppressed at L →∞ unless α L = 1 2 .

Below we provide a couple of short comments about this result. The proof and full DMFT is provided in Appendix E.4.

At initialization t = 0 , the only term which contributes to the residual stream layer dynamics is the integrated Brownian motion ∫ τ 0 du ( τ ′ ) which survives at infinite depth for α L = 1 2 . For α L = 1 this term disappears in the limit. The structure of C ( τ ) is also modified by additional response functions at α L = 1 2 [11] which we show disappear for α L = 1 .

The weights within residual blocks (including the MHSA block) can be treated as completely frozen for α L < 1 in the L →∞ limit, which leads to the simplified statistical description of the preactivations in those layers. However, the residual stream variables h ( τ ) do still obtain Θ L (1) updates. At α L = 1 the weights in the MHSA blocks evolve by Θ L (1) , causing additional feature evolution in the model.

A consequence of our large N and large L result is that the N,L →∞ limit with α L = 1 2 (the parameterization studied by [11, 12]) would lead to A ℓ hss ′ ( x , t ) = 0 for all time t . Thus the MHSA blocks would only involve average pooling operations over the spatial indices, despite the residual stream kernels H ℓ updating from feature learning.

10 1

10 2

L

10

1

10 0

|

Wq

,

k

(

t

)

Wq

,

k

(0)|

L

1/2

L

= 1

2

L

=1

(a) Key/Query Weight Changes

10 5

10 6

10 7

10 8

10 9

10 10

Compute

10 0

6×10 1

Test Loss

L

= 1

2

L

=1

L

=4

L

=8

L

=16

L

=32

(b) Depth Scaling for

α

L

∈ {

1

2

,

1

}

First, we note in Figure 5 that the weights within each attention block freeze as L →∞ with α L = 1 2 case but move at a constant scale for α L = 1 . As a consequence, the loss at large L can be lower in the α = 1 parameterization.

We can see some numerical evidence for the first of these effects in Figure 6 (a)-(b) where initially training at large L is slower than the base model and the initial kernel appears quite different for L = 4 and L = 64 . The initial kernel will decrease in scale as L → ∞ for α L = 1 since the preactivation vectors lose variance as we discuss in Appendix E.4, resulting in slower initial training. However, we note that the final learned feature kernels are quite similar after enough training.

In summary, our results indicate that the α L = 1 parameterization is the one that allows attention layers to actually be learned in the limit L → ∞ , but that this parameterization leads to a less structured kernel at initialization.

4 Experiments in Realistic Settings

In practice, large scale neural networks do not generally operate close to their limit. Given the costs of training large networks, one would ideally operate in a regime where there is a guarantee of consistent improvements with respect to model scale. In pursuit of this goal, we apply our theoretical findings of this paper to training language models on a larger natural language dataset, a Transformer with causal attention blocks trained on the C4 dataset [43] with Adam optimizer. As mentioned in 2.2, while our exact theoretical description of these infinite limits focus on SGD, we can implement an appropriate scaling for Adam which preserves the scale of internal feature updates. This allows us

10 2

10 3

10 4

10 5

Steps

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

2.25

Test Loss

Base

=128

N

=128

L

=64

(a) Training Dynamics

Initial

H L

(

x

,

x

′

)

Base

Final

H L

(

x

,

x

′

)

=128

N

=128

L

=64

(b) Initial and Final Residual Stream Pooled Kernels

Initial

H L

′

Base

Final

H L

′

=128

N

=128

L

=64

1

0

1

0

1

2

3

4

5

Initial Density

= 1

2

=1

15

10

5

0

0.0

0.2

0.4

0.6

0.8

Final Density

N

=16

N

=32

N

=64

N

=128

(d) Attention Distributions Before and After Training

to investigate realistic training dynamics of our LLM as we take the N,L, H → ∞ limits. Training details are provided in Appendix F

In Figure 7 (a), we sweep over each of the model dimensions independently for each parameterization of α A ∈ { 1 , 1 2 } on the left and right respectively. For fixed N and L , scaling H provides a similar increase in performance in both parameterization and appear to start converging to a final loss around 5 , with slight benefit to α A = 1 2 . For fixed H and L , scaling N provides a similar increase in performance to scaling heads in when α A = 1 , but a substantial increase when α A = 1 2 . This is in line with our predictions in Section 3.2 about the benefits of diversity across attention heads. Next, for fixed N and H , scaling L provides little to no benefit in either parameterization as predicted in Section 3.4. Finally, we inspect the sample and spatial residual stream kernels of these models before and after training and find that the kernels are identical for both α A , except for a slight difference for large N . Furthermore, they are extremely similar for large N and large H .

Taken together, these results suggest that scaling different model dimensions do indeed have different effects on training dynamics and final performance. This provides groundwork for future large-scale experiments systematically investigating their trade-offs, thereby identifying compute-optimal scaling of realistic architectures in parameterizations with well-defined limits.

5 Discussion

This paper provided analysis of the infinite head, depth and key/query dimension limits of transformer training in the feature learning regime. We showed that feature learning in µ P multi-head transformers in the limit of N →∞ collapses to single-head self-attention. At finite N and infinite heads H → ∞

0

2000

4000

6000

8000

10000

3

4

5

6

7

8

9

10

11

Loss

α

A

= 1

Base

H

= 16

H

= 32

H

= 64

H

= 128

0

2000

4000

6000

8000

10000

3

4

5

6

7

8

9

10

11

α

A

= 1

2

Base

H

= 16

H

= 32

H

= 64

H

= 128

0

2000

4000

6000

8000

10000

3

4

5

6

7

8

9

10

11

Loss

Base

N

= 16

N

= 32

N

= 64

N

= 128

0

2000

4000

6000

8000

10000

3

4

5

6

7

8

9

10

11

Base

N

= 16

N

= 32

N

= 64

N

= 128

0

2000

4000

6000

8000

10000

Steps

3

4

5

6

7

8

9

10

11

Loss

Base

L

= 8

L

= 16

L

= 32

L

= 64

0

2000

4000

6000

8000

10000

Steps

3

4

5

6

7

8

9

10

11

Base

L

= 8

L

= 16

L

= 32

L

= 64

(a) Training dynamics of LLM on C4

Initial

H L

(

x, x

′

)

Base

Final

H L

(

x, x

′

)

:

α

A

= 1

Final

H L

(

x, x

′

)

:

α

A

= 1

2

H

= 128

N

= 128

L

= 64

(b) Kernels: final token across samples

Initial

H L

ss

′

Base

Final

H L

ss

′

:

α

A

= 1

Final

H L

ss

′

:

α

A

= 1

2

H

= 128

N

= 128

L

= 64

(c) Kernels: tokens within single sample

we showed that there is an alternative limit which maintains a distribution over attention heads . We discussed two different large depth limits of transformer training that reduce to differential equations in the residual layer time τ . The depth scaling that maintains feature learning within all MHSA blocks ( α L = 1 ) causes the initial kernel to lose structure from the initialization as L → ∞ , but allows learning of the self-attention variables, whereas the depth scaling that preserves structure from initialization ( α L = 1 2 ) leads to static layers.

Limitations and Future Directions Currently exact theoretical analysis of the limit is focused on SGD (and can be easily extended to SGD+momentum [15]) while Adam is currently only reasoned with rough scaling arguments rather than an exact theoretical description of the limit. Since Adam is most commonly used to train transformers, a theory of the limiting dynamics of Adam in Transformers would be an important future extension. In addition, while we provide an exact asymptotic description of network training, the limiting equations are compute intensive for realistic settings which is why we focus our empirical investigations on training large width networks in the appropriate parameterizations. Lastly our techniques assume that the number of training steps is fixed as the scaling parameters of interest ( N, H , L ) are taken to infinity. However, it would be important to understand learning dynamics in the regime where model size and training times are chosen to balance a compute optimal tradeoff (or perhaps even training longer than compute optimal) [8, 39, 44]. In this regime, harmful finite model-size effects become significant and comparable to the finite training horizon [10, 16, 17, 39]. Thus stress testing the ideas in this work at larger scales and longer training runs would be an important future direction of research into scaling transformer models.

References

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Kai Han, Yunhe Wang, Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, et al. A survey on vision transformer. IEEE transactions on pattern analysis and machine intelligence , 45(1):87-110, 2022.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab, Matthias Minderer, and Yi Tay. Scenic: A jax library for computer vision research and beyond. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 21393-21398, 2022.

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.

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.

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.

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.

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems , 34:17084-17097, 2021.

Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview. net/forum?id=KZJehvRKGD .

Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Feature learning in infinite depth neural networks. In The Twelfth International Conference on Learning Representations , 2023.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

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.

Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems , 35:32240-32256, 2022.

Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales, 2023.

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

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization, 2023.

Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning , pages 2793-2803. PMLR, 2021.

Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems , 35:27198-27211, 2022.

Bobby He and Thomas Hofmann. Simplifying transformer blocks. arXiv preprint arXiv:2311.01906 , 2023.

Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli. Geometric dynamics of signal propagation predict trainability of transformers, 2024.

Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and-width limit. Advances in Neural Information Processing Systems , 36, 2024.

Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/ forum?id=RbLsYz1Az9 .

Nicola Muca Cirone, Maud Lemercier, and Cristopher Salvi. Neural signature kernels as infinite-width-depth-limits of controlled resnets. arXiv preprint arXiv:2303.17671 , 2023.

Lénaïc Chizat and Praneeth Netrapalli. The feature speed formula: a flexible approach to scale hyper-parameters of deep neural networks, 2024. URL https://arxiv.org/abs/2311. 18718 .

Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming-Yu Liu. On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems , 33:21370-21381, 2020.

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813 , 2023.

Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic gradient descent: Deep learning without hyperparameters. arXiv preprint arXiv:2304.05187 , 2023.

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.

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.

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

Haim Sompolinsky and Annette Zippelius. Dynamic theory of the spin-glass phase. Physical Review Letters , 47(5):359, 1981.

Moritz Helias and David Dahmen. Statistical field theory for neural networks , volume 970. Springer, 2020.

Stefano Sarao Mannelli, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborova. Passed &spurious: Descent algorithms and local minima in spiked matrix-tensor models. In international conference on machine learning , pages 4333-4342. PMLR, 2019.

Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems , 33:9540-9550, 2020.

Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zdeborova. Rigorous dynamical mean field theory for stochastic gradient descent methods. arXiv preprint arXiv:2210.06591 , 2022.

Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws, 2024.

Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A , 8(1):423, 1973.

Blake Bordelon and Cengiz Pehlevan. The influence of learning rule on representation dynamics in wide neural networks. arXiv preprint arXiv:2210.02157 , 2022.

Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. arXiv preprint arXiv:2010.08127 , 2020.

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. Journal of machine learning research , 21(140):1-67, 2020.

Ibrahim M Alabdulmohsin, Xiaohua Zhai, Alexander Kolesnikov, and Lucas Beyer. Getting vit in shape: Scaling laws for compute-optimal model design. Advances in Neural Information Processing Systems , 36, 2024.

Etai Littwin and Greg Yang. Adaptive optimization in the ∞ -width limit. In The Eleventh International Conference on Learning Representations , 2022.

Mufan Li, Mihai Nica, and Dan Roy. The neural covariance sde: Shaped infinite depth-andwidth networks at initialization. Advances in Neural Information Processing Systems , 35: 10795-10808, 2022.

Appendix

A Additional Figures

10 0

10 1

10 2

10 3

Steps

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Test Loss

=16

=32

=64

=128

(a) Vary

H

10 0

10 1

10 2

10 3

Steps

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Test Loss

N

=16

N

=32

N

=64

N

=128

(b) Vary

N

10 0

10 1

10 2

10 3

Steps

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Test Loss

L

=4

L

=8

L

=16

L

=32

L

=64

(c) Vary

L

Initial

H L

(

x

,

x

′

)

=16

Final

H L

(

x

,

x

′

)

=32

=64

=128

Initial

H L

(

x

,

x

′

)

N

=16

Final

H L

(

x

,

x

′

)

N

=32

N

=64

N

=128

Initial

H L

(

x

,

x

′

)

L

=4

Final

H L

(

x

,

x

′

)

L

=8

L

=16

L

=32

限

Initial

H L

′

=16

Final

H L

′

=32

=64

=128

Initial

H L

′

N

=16

Final

H L

′

N

=32

N

=64

N

=128

Initial

H L

′

L

=4

Final

H L

′

L

=8

L

=16

L

=32

0

500

1000

1500

2000

2500

Steps

1.4

1.6

1.8

2.0

2.2

Test Loss

N

=8

N

=16

N

=32

N

=64

N

=128

N

=256

N

=512

(a)

N

Scaling with

α

= 1

0

500

1000

1500

2000

2500

Steps

1.4

1.6

1.8

2.0

2.2

Test Loss

N

=8

N

=16

N

=32

N

=64

N

=128

N

=256

N

=512

(b)

N

Scaling with

α

=

1

2

10 9

10 10

10 11

10 12

10 13

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

Loss

α

A

= 1

Base

H

= 16

H

= 32

H

= 64

H

= 128

10 9

10 10

10 11

10 12

10 13

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

α

A

= 1

2

Base

H

= 16

H

= 32

H

= 64

H

= 128

10 9

10 10

10 11

10 12

10 13

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

Loss

Base

N

= 16

N

= 32

N

= 64

N

= 128

10 9

10 10

10 11

10 12

10 13

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

Base

N

= 16

N

= 32

N

= 64

N

= 128

10 9

10 10

10 11

10 12

10 13

Compute

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

Loss

Base

L

= 8

L

= 16

L

= 32

L

= 64

10 9

10 10

10 11

10 12

10 13

Compute

4

.

0

4

.

5

5

.

0

5

.

5

6

.

0

6

.

5

7

.

0

Base

L

= 8

L

= 16

L

= 32

L

= 64

B Implementations for Vision and Causal Language Modeling Transformers

We provide an example FLAX implementation of the vision transformer and causal language model.

We start by defining a fixed layernorm operation

The MHSA layer is implemented as the following where scale_exp represents α A .

The two layer MLP block is implemented as the following with ϕ = gelu nonlinearity.

We also allow for a trainable positional encoding matrix.

Each residual block is implemented as the following. Below we show the α L = 1 implementation.

Our vision transformer model consists of an embedding layer which is applied to each patch, a positional encoding layer, L residual layers each containing a MHSA and MLP block, a spatial pooling operation, and a readout.

For the causal decoder only model, we need to modify the Attention layer and also prevent pooling over spatial indices before the readout.

C Simple Heuristic Scaling Analysis

In this section, we heuristically work out the simple scaling analysis to justify the set of parameterizations and learning rates we consider. More detailed theoretical analysis for the limit of SGD training is provided in Appendix E where we exactly characterize the N →∞ , H → ∞ and L →∞ limits. We consider taking heads H , inner dimension N and depth L to infinity separately and attempt to control the scale of gradients and updates.

C.1 Learning Rate Scalings

We show that the correct learning rate scaling for SGD is η = η 0 N H L 2 α L -1 . For Adam, the learning rate should be scaled as η = η 0 N -1 / 2 H -1 / 2 L -1+ α L .

C.2 Heuristic Analysis of Feature Changes Under SGD

In this section we consider performing a single update on a single example to all weight matrices.

where g ℓ +1 ∈ R N H and v ℓ h ∈ R N have Θ(1) entries. Thus, computing a perturbation to the forward pass we find

The term in the brackets is Θ(1) and we see that the perturbation from each layer contributes Θ( L -1 ) . As there are L layers, this will give a total change to the final layer h L that is Θ(1) .

For the Attention variables, we note that the

where q h , k ℓ h ∈ R N are the query and key for head h and h ∈ R N H is the residual stream preactivation. We can thus compute the changes to the keys and queries due to changes in their associated weights

where H ℓ = 1 N H h ℓ · h ℓ ∼ Θ(1) . Combining these changes , we find the following update to the pre-Attention variables A ℓ h = 1 N α A k ℓ h · q ℓ h

since 1 N k · k , 1 N q · q ∼ Θ(1) . This update to the attention variable due to changes in W ℓ K , W ℓ Q will clearly die out as L →∞ unless α L = 1 .

C.3 Heuristic Analysis of Feature Changes Under Adam

For Adam, the gradient of each individual parameter entry is approximately normalized by its scale [45]. Thus the learning rate η sets the size of the updates. This is why we scale the learning rate as η = 1 L 1 -α L √ N H which gives the same scale updates to the weights as SGD

Again computing the correction to the forward pass we find

Similarly our update generates the same scale of weight updates to the key and query weight matrices

We can therefore follow the identical argument to identify the scale of the change to the pre-attention variables δ A ℓ h = Θ( L 1 -α L ) .

C.4 What Counts as Feature Learning for Attention Layers?

Any parameterization with α N ∈ [ 1 2 , 1] will cause all updates to A ℓ h and entries of h ℓ +1 to be Θ N,H,L (1) across finite N . The entries of q and k only move by Θ N (1) if α A = 1 ( µ P scaling). However, we argue that this criterion is not strictly necessary. Rather, feature learning could alternatively be defined in terms of evolution of macroscopic variables ( H , A , f , etc) rather than preactivation or key/query vector entries themselves. Table 3 summarizes two example values of α A which are of special interest for their N →∞ limits.

2

The choice α A = 1 2 allows the variance of A ℓ h to be constant size as a function of N while also enabling learning of these variables. We verify these scalings in Figure 13.

10 2

10 3

10 4

N

10

10

10

8

10

6

10

4

10

2

10 0

10 2

(

k

)

2

t

=1

t

=200

= 1

2

=1

N

1

N

2

(a) Change in

k

Entries (SGD)

10 2

10 3

10 4

N

10

12

10

9

10

6

10

3

10 0

10 3

(

)

2

t

=1

t

=200

= 1

2

=1

N

2

(b) Change in

A

Entries (SGD)

10 2

10 3

10 4

N

10

7

10

6

10

5

10

4

10

3

10

2

10

1

10 0

(

k

)

2

t

=1

t

=200

= 1

2

=1

N

1

(c) Change in

k

Entries (Adam)

10 2

10 3

10 4

N

10

5

10

4

10

3

10

2

10

1

10 0

10 1

(

)

2

t

=1

t

=200

= 1

2

=1

(d) Change in

A

Entries (Adam)

D DMFT Primer and Simple Examples

D.1 Main Conceptual Idea of the Approach

Dynamical mean field theory is a method that was developed in the physics of spin glasses for dealing with dynamical systems that depend on a fixed source of disorder . The disorder could be random couplings between sites in a spin glass model [34], random connections between neurons in a random recurrent neural network [35], random data drawn from a distribution [37, 39] or the random initial weights in a deep neural network [15, 11]. In our case, we are interested in the last example, where the feature learning dynamics of a randomly initialized transformer is a function of the initial weights in each layer. In what follows, we will give a primer on the main objects which typically appear in a DMFT analysis (the correlation and response functions) to illustrate the main ideas of the approach.

D.2 Example 1: Linear Dynamics with GOE Matrix

In this section, we discuss and derive the DMFT equations for the simplest possible example, a linear dynamical system with a Gaussian symmetric coupling matrix.

In this example we show that the DMFT path integral is computing something non-trivial about the kinds of dynamics induced by a linear dynamical system with a random matrix. In this linear example, the DMFT path integral encodes spectral properties of the random matrix.

Let's consider the simplest possible example: d dt h i ( t ) = 1 √ N ∑ N j =1 W ij h j ( t ) where W ij = W ji is a Gaussian symmetric matrix (GOE). This matrix is fixed while the state h ( t ) ∈ R N evolves. The path integral appraoch would tell you that in the N →∞ limit, every neuron i has identical statistics given by the stochastic integro-differential equation

where ⟨·⟩ denotes an average over the random variables u ( t ) . In this picture, the averages ⟨⟩ over the noise can also be interpreted as averages over all N neurons in the system, each of which are independent. This stochastic equation can be used to close the evolution equations for the correlation C ( t, s ) and linear response function R ( t, s ) .

A generic result of this path integral DMFT picture is

All neurons (all variables h i ) decouple statistically. The presence of all other neurons only enters through "macroscopic" quantities C ( t, s ) and R ( t, s ) known as the correlation and response functions. The distribution of these functions over random realizations satisfies a large deviations principle where the distribution over C, R has the form p ( C, R ) ∼ e -NS ( C,R ) where S is the DMFT action obtained from the path integral method.

Extra memory terms like ∫ t 0 R ( t, s ) h ( s ) appear which depend on the state at earlier times s < t . The Markovian (deterministic) system for p ( h | W ) becomes stochastic and nonmarkovian after marginalizing p ( h ) = ∫ d W p ( h | W ) p ( W ) . I would argue these memory terms are not obvious apriori but are systematic to compute in this framework.

Since this toy example is a linear dynamical system , one can also identify a connection between the DMFT correlation and response and spectral properties of the random matrix W . We note that the response has the form

where ρ ( λ ) is the eigenvalue density of W . In fact a Fourier transform of our DMFT equation recovers the semicircle law ρ ( λ ) = 1 π Im R ( iλ ) = 1 2 π √ [4 -λ 2 ] + for the eigenvalues.

In general, one can think of DMFT as a more powerful version of this method that can also handle nonlinearities.

D.3 Example 2: Deep Linear Network Updates

In this section I will try showing how this DMFT approach can give useful insights into reasoning about learning updates which are not obvious apriori. While our paper advocates for taking depth L →∞ in a residual network, we first thought about simply scaling depth in a standard MLP. Below we show how the proliferation of response terms gives a different predicted scaling with L than if we naively disregarded response terms.

Consider a non-residual linear MLP network with µ P/mean-field scaling with L hidden layers with N →∞ . Train the model for a single step of gradient descent with learning rate η on a data point ( x, y ) with | x | 2 = 1 and y = 1 and output multiplier 1 /γ 0 . The forward pass variables h ℓ ( t ) and the backward pass variables g ℓ ( t ) are defined recursively as

Now, naively, one may think that 1 √ N W ℓ (0) h ℓ ( t ) has entries that are independently Gaussian with covariance H ℓ ( t, t ′ ) . However, this is incorrect and the DMFTresponse functions give an additional correction.

The DMFT Equations for this Model Following the approach of [15, 11], we find the following DMFT equations for the preactivations h ℓ after 1 step of training

where the random variables u ℓ (0) , u ℓ (1) and r ℓ (0) have the following covariance structure

and the feature kernels H ℓ ( t, t ′ ) , G ℓ ( t, t ′ ) and response functions A ℓ ( t, t ′ ) have the form

These recursions can be solved with the initial conditions H ℓ (0 , 0) = 1 and G ℓ (0 , 0) = 1 which implies that H ℓ (1 , 0) = 1 so that

Using this equation, we find

This is the DMFT prediction for the scale of the feature kernels after a step of training.

Neglecting the DMFT Response Gives Incorrect Depth Scalings for MLPs However, if we had neglected the DMFT response functions and approximated the dynamics as

The feature variance after t = 1 step of gradient descent H ℓ = 〈 h ℓ (1) 2 〉 after t = 1 step, the final layer

We see that without the response terms we get a completely different scaling prediction with L !

Scaled Residual Networks For 1 √ L residual block scaling ( α L = 1 2 ), the response functions are still important to accurately characterize the dynamics and contribute Θ L (1) corrections to the feature learning dynamics as L →∞ . However, for the 1 /L block multiplier scaling, the response functions do not contribute in the limit. These facts are not-apriori obvious but follow from the DMFT analysis (either path integral or cavity approach).

E DMFT Analysis for Transformers

In this section, we derive the limiting equations for the infinite head limit H → ∞ of training with SGD. The results can be easily extended to SGD with momentum following the methods of [11, 15].

E.1 Deriving the DMFT Action

In this section we will derive the limiting equations of motion for stochastic gradient descent in the H → ∞ limit. We start by defining the loss function which we aim to minimize

where p ( x ) is the data distribution of interest. We note that this can be the population loss or the empirical loss on a finite collection of points. We let ∆( x, t ) ≡ -∂ℓ [ f ( x ,t )] ∂f ( x ,t ) represent the error signal on datapoint x . At each step of training t a batch of examples B t = { x 1 ( t ) , ..., x | B t | ( t ) } is generated and used to estimate a gradient for SGD. In what follows, we let E x ∼ B t represent averages over the minibatch at time t . We emphasize here that the batches B t are assumed to be given or fixed and are not averaged over as random draws from p ( x ) , but rather our expectation simply denotes the empirical mean over the minibatch at time t

We start by expressing again the forward pass equations for each layer. To make this analysis more compressed while still capturing all of the interesting aspects, we will first compute the equations of motion in the absence of MLP layers (which were analyzed in prior works, see Appendix E) and layernorm (which in the limit it will only apply a deterministic affine transformation to each of the entries of the residual stream and the backward pass gradient variables as we will show explicitly in Appendix E.6). In the absence of layernorm, our forward pass has the form

where the MHSA layer is

To compute the weight dynamics we again introduce the necessary gradient fields which we previously argued have Θ(1) entries

We also introduce the following intermediate quantities which are necessary to characterize the backward pass through the attention layer

We need to break up each of the weight matrices into their initial component and their update from SGD

We can now express the residual stream as

where we introduced the fields

We see that, regardless of the choice of α L , the update to the residual stream due to feature learning will scale as 1 /L which is necessary for a stable infinite depth limit [11]. To track the dynamics of the value vectors, we must also track the variables

We similarly find the following for the key dynamics

and an analogous update equation holds for the query dynamics q

We see that if the residual stream is frozen so that χ K , χ Q are static, the keys and queries will only evolve in the N,L →∞ limits if α L = α A = 1 as we argued in the main text. From these equations we can deduce the pre-attention values A ℓ hss ′ ( x , t ) .

Next, we examine the dynamics of the M ℓ hs ( x , t ) variables which are defined as

We can verify that M ℓ h are all Θ N, H (1) throughout training by expanding the dynamics of the attention output o ℓ hs ′ ( x , t ) .

E.1.1 Backward Pass

Next, we need to work out the recursions for the backward pass variables. After these equations have been worked out, we can isolate the dependence of the full dynamics on all of the initial weights.

MHSA Layer We will start by differentiating through the MHSA layer

where we introduced the variables

and their associated kernels

where we introduced the following random fields

E.1.2 Why we need α A = 1 for gradient stability

From the previous equation we see that regardless of the choice of α A we need to choose the variance of W ℓ K h (0) and W ℓ Q h (0) so that 1 √ N W ℓ K h (0) q ℓ h ( x , t ) has O (1) entries. This means that the entries can at most Θ(1) and not Θ( N 1 -α A ) as we originally stipulated in order to obtain k and q with Θ(1) entries at initialization. Thus, with this required scaling, we must have either α A = 1 and Θ(1) variance of the weights, which leads to attention variables which are Θ( N -1 / 2 ) at initialization or we choose α A = 1 2 and choose k , q to have entries of scale Θ( N -1 / 2 ) at initialization. These both lead to the same vanishing initial condition for the pre-attention variables. It thus suffices to consider µ P scaling α A = 1 to study the N →∞ limit. We stress that this effect is not visible from a simple heuristic analysis of the forward pass variables after an update like we perform in Appendix C.

Isolating all Dependence on Initial Conditions To summarize the previous sections, we begin by collecting all of the stochastic fields which show up in the dynamics and depend on the initial weight matrices. These quantities all come in pairs since for each matrix we need to consider the forward and backward passes through the initial matrix.

The following variables are necessary to characterize the dynamics of the N H -dimensional residual stream

From these variables, we can reconstruct the entire dynamics for the h , g fields. We therefore study the moment generating functional for the above primitive random variables. Let

We multiply by the identity to enforce the definition of these random variables in terms of the initial weights. As an example, we would have for the first random variable χ ℓ +1 s ( x , t ) and its corresponding pair ξ ℓ s ( x , t )

Attention Output Matrices In this section we integrate over W ℓ O h , W ℓ K h , W ℓ Q h .

where we introduced the response function

Value Matrices Next, we average over the value matrices W ℓ V h (0) which gives

Key/Query Matrices Next, we need to perform the averages involving W K h (0) which has entries resulting in

We follow an identical procedure for the query matrices W ℓ Q (0) .

Enforce Kernel Definitions After this step, we can introduce new resolutions of the identity for each of the kernels that appeared in the above computation

This is repeated for all of the response functions which involve sums over N H variables

There are other kernels which are only relevant within a single head including { Q ℓ h , K ℓ h , A ℓ h , M ℓ h } . These

We now combine all of the order parameters into a large collection Q which are vectorized over all layer, time, spatial and sample indices

After introducing this collection of order parameters, our original MGF satisfies a large deviation principle with

where the DMFT action S ( Q ) has the form

The residual stream single site moment generating functional has the form

Though this expression is cumbersome, we will show that, since ˆ H, ˆ G vanish at their saddle point this MGF merely encodes the following statistical description of the fields of interest such as

We note that h only depends on ¯ χ ℓ O = 1 √ H ∑ H h =1 χ ℓ O h σ ℓ h so it has the form

This fact will be important when we take the large head limit with N fixed in Appendix E.3. Next, we analyze the single site MGF for the hidden MHSA layers

E.2 Infinite N (Key/Query dimension) Limit

First, we take the N →∞ limit with H , L fixed. This can be obtained with a simple saddle point procedure using the action in the form written in the previous section. This calculation exactly mimics prior works [15, 11] where all of the order parameters Q take on their values at the saddle point Q ⋆ .

Saddle Point Values for Order Parameters Under this saddle point, all of the order parameters presented will concentrate and all neurons will become statistically independent. The governing equations for the order parameters Q ⋆ are given in terms of averages ⟨·⟩ over the single site densities defined by the moment generating functionals Z and have the form

Single Site Stochastic Processes Our fields of interest will obey the stochastic dynamics

Residual Stream The forward pass residual variables obey the following stochastic process

where

the backward pass satisfies

The keys and queries have dynamics

E.2.1 Multi-head Attention is Single-Head Attention as N →∞

In this section, we use the derived saddle point equations for the N →∞ limit and argue that they imply that all heads in the MHSA layer learn identical attention matrices and contribute the same feature updates to the residual stream. To do so, we proceed in a three step inductive argument.

First, we show that at initialization, all key, query, value and attention matrices { Q h , K h , V h , A h , M h } are equal across heads.

Next, we show inductively that if these quantities { Q h , K h , V h , A h , M h } are identical across heads up to some time, then that implies that the response functions R h are also identical across heads up to that time.

Lastly, we show that if the response functions R h are identical up to some time, then that implies that the MHSA kernels { Q h , K h , V h , A h , M h } will also be identical across heads at future times.

First, we note that, at initialization, all of the MHSA kernels are identical across heads since

Next, we need to analyze the response functions under an inductive hypothesis on the equality of the MHSA kernels. We start by noting that all response functions are causal so we can group the response functions that arise from χ ℓ O h with the feature learning update to the residual stream, writing the following compressed equation for the forward and backward passes

where C h ℓ and C g ℓ only involve deterministic head-averaged kernels and thus do not carry a h index. We now derive useful response function identities

These equations give the needed response function R g ℓ +1 χ ℓ O . From these equations we immediately see that if σ ℓ hss ′ ( x , t ) = σ ℓ h ′ ss ′ ( x , t ) (which holds under our inductive hypothesis) then R g ℓ +1 χ ℓ O h = R g ℓ +1 χ ℓ O h ′ . This same argument is repeated for all other response functions that are computed as derivatives of residual stream variables. We have thus found that

Now, we can analyze the dynamics of the keys, queries and values within a head using the above property and the original inductive hypothesis that A h = A h ′ , M h = M h ′ ... for times less than t . We will now prove that this implies that these variables will remain the same. We start by examining the keys and queries which have the form

where C k ℓ ss ′ ( x , x ′ , t, t ′ ) and C q ℓ ss ′ ( x , x ′ , t, t ′ ) are two operators that do not carry a h (are the same across all heads). These relations can be viewed as a linear system of equations. We let k ℓ h = Vec { k ℓ hs ( x , t ) } s x t and analogously C k ℓ = Mat { C k ℓ ss ′ ( x , x ′ , t, t ′ ) } as in the calculations of Bordelon and Pehlevan [15, 41]. Using this shorthand, we can express the key/query and attention kernels as

We thus see that the final kernels K ℓ h , Q ℓ h , A ℓ h are all identical across heads. An identical argument can be carried out for the value kernel V ℓ h and the M ℓ h order parameter.

E.3 Infinite H Limit

In this section, we compute the infinite head limit with N,L fixed. This limit is more technically involved than the N →∞ limit which required only a simple saddle point of the full DMFT action over all kernels. At finite N we cannot use this technique since the kernels within MHSA blocks are random variables . However, as was shown in Bordelon and Pehlevan [17], the DMFT action still contains the necessary information to characterize the distribution over order parameters at finite N . In the case of transformers with infinitely many heads, a subset of the order parameters introduced in

the previous section will still concentrate as H → ∞ including

and many more correlation and response functions. We will call the full collection of all the necessary head-averaged order parameters Q global. Further, not all of the stochastic fields will be relevant to characterize the residual stream. Specifically, only head-averaged fields ¯ χ ℓ O , ¯ ξ ℓ V , ¯ ξ ℓ K , ¯ ξ ℓ Q are relevant. For example, the first of these is defined as

We note that the above equation is true at any value of N but the covariance of u ℓ h and the response functions R v ℓ ξ ℓ O h are random variables [17]. However, the residual stream only depends upon collective, head-averaged variables

The intuition behind the large H limit is that, even though σ ℓ h and R v ℓ ξ ℓ O h are random variables, there should be a central limit theorem for ¯ u ℓ O and a law of large numbers for ¯ R v ℓ ξ ℓ O σ ss ′ ( x , x ′ , t, t ′ ) .

E.3.1 Partitioning Order Parameters

Based on the intuition developed in the previous section, we now derive an alternative DMFT action by tracking the moment generating functional for the head-averaged random fields that occur on the residual stream ¯ χ ℓ O , ¯ ξ ℓ V , ¯ ξ ℓ K , ¯ ξ ℓ Q . To characterize this joint distribution, we must also of course keep track of the random fields within the MHSA blocks such as { χ ℓ Q h , χ ℓ K h , χ ℓ V h , ξ ℓ O h } h ∈ [ H ] . Repeating the path integral setup of the previous section, we need to performing averages over all initial weights such as

It is clear that from this integral that the relevant self-averaging order parameters are

We repeat this same procedure for all collections of weights and arrive at the following set of order parameters

We expect that these order parameters will be self-averaging since they involve averages over H variables. However, the other variables we introduced {A h , Q h , K h , V h } will not concentrate at finite N and will instead behave as random variables.

After introducing these order parameters we find that the moment generating functional has the form

where S has the form

The single site moment generating function for the residual stream has the form

We can express the MHSA single-head partition functions Z MHSA in terms of the remaining order parameters within each head that will no longer concentrate at finite N

After introducing these order parameters, we have

where the key-query-value single site partition function has the form

The saddle point equations for this limit are computed as derivatives with respect to Q global only , reflecting that head-averages will converge as H → ∞

The final saddle point equations are given in terms of averages over the distribution of heads defined by Z MHSA which we denote as ⟨·⟩ MHSA as well as averages over the residual stream which we denote as ⟨·⟩ .

These equations give the following (we suppress the sequence indices to simplify the final expressions)

Residual Stream Dynamics The residual stream satisfies the following single-site dynamics

This matches the result provided in the main text which introduces a compressed

where p t ( x ) = 1 | B t | ∑ x ′ ∈ B t δ ( x -x ′ ) denotes the uniform distribution over the batch B t .

E.4 Infinite L Limits

In this section, we discuss the two large L limits. This can be derived formally in two distinct ways. First, one could start with the initial For this section, it suffices to reason about the scale of the Gaussian noise which appears in the residual stream and the contribution from the response functions.

E.4.1 Basic Intuition

Deriving the infinite depth limits To gain intuition for the large H , L →∞ limit, we use the fact that the random variables χ ℓ = u ℓ + R ℓ g ℓ decompose into a Gaussian u ℓ which are uncorrelated across layers and a linear response R ℓ are response functions. This implies that

We first note that the sum of the Gaussians is a zero-mean random variable with standard deviation

Thus, this integrated random variable will vanish unless α L = 1 2 .

Next, we can investigate the scale of the residual stream response functions. For instance

As a consequence, we see that the effect of the Gaussian and linear response terms will vanish as L →∞ provided that α L > 1 2 . We will consider first, the case where α = 1 which gives an ODE like limit for the residual updates before moving onto the more involved α L = 1 2 case.

To formally take the L →∞ limit, we redefine all of the preactivation fields and kernels in terms of layer time τ defined as

For example, the residual kernels are defined as

The finite difference equations for the residual updates L ( h ℓ +1 -h ℓ ) ∼ ∂ ∂τ h ( τ ) become differential updates (either SDE-like or ODE-like depending on α L ) [46, 11].

E.4.2 ODE Limit α L = 1

First, we investigate the case of α L = 1 . In this case, the H , L → ∞ limit results in a complete disappearance of the ¯ χ ℓ O , ¯ ξ ℓ V , ¯ ξ ℓ K , ¯ ξ ℓ Q fields.

E.4.3 SDE Limit α L = 1 2

This α L = 1 2 limit is more technically involved since neither the Gaussian terms from the DMFT nor the response functions vanish. For the Gaussian terms, we note that the sums of the independent Gaussians are all multiplied by 1 √ L ∼ √ dτ , which can be interpreted as integrated Brownian motion in the limit

Following the derivation of Bordelon et al. [11], which maintains the exact dependence on the full integrated response and provides the result as an integrated SDE

Combining the forward pass equations from the previous two subsection recovers the Result 3 of the main text. This is combined with a complementary equation for the backward pass.

E.5 Effect of MLP Layers

Adding the MLP block to the residual stream can also be easily handled with the methods of the preceeding sections. The forward pass equations in this case take the form

where the MLP layer is

The following gradient fields are necessary

The MLP hidden layer dynamics is much simpler to characterize and resembles the structure analyzed in prior works on infinite width networks [15].

Again, we see that the inner dynamics for ˜ h ℓ, 1 due to the weight updates in this layer scale as L -1+ α L , suggesting the need to choose α L = 1 if we desire this hidden layer to contribute to the representational updates.

MLP Layer Gradients For the MLP layer we have the simpler backpropagation equations

The components of these fields that depend on initial conditions are

MLP Matrices After utilizing these resolutions of the identity for all s , x , t , we can integrate over the weights W ℓ, 2 (0)

where we introduced the response function

We can perform an identical step to integrate over W ℓ, 1 (0) . This gives us

where we introduced

E.6 Effect of Layer Norm on the Limiting Process

Layernorm The derivative of layer-norm ∂ ¯ h ℓ s ′ ∂ h ℓ ⊤ s ′ acts as the following in the large H limit

In the limit of H → ∞ the variables µ = 1 N H h · 1 and σ 2 = 1 N H | h -µ 1 | 2 will become deterministic over random initializations. We thus just have to consider how these types of vectors act on gradients

Each of these operations will lead to one inner product that will be self averaging µ g = 1 N H 1 · g or ( 1 N H [ h -µ 1 ] ⊤ g ) . Thus this operation will not alter the backward pass in terms of scaling.

F Compute Resources and Experimental Details

Each of the experimental runs performed in this paper were all performed on single NVIDIA H100 GPU. Each run of the full CIFAR-5M took anywhere from 5 minutes to 1 hour depending on model size. Each run of the C4 training took anywhere from 1 hour to 6 hours depending on model size and total amount of training steps.

The language model trained on C4 used a context length of 256 and the GPT-tokenizer from Huggingface. Sequences that were too short were concatenated with other sentences to reach the full context length rather than padding the end of the sequence. We use trainable positional encodings and separate embedding and decoding parameters as implemented in Appendix F.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We claim to analyze various limits of transformer training and provide theoretical and empirical results to that effect.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We provide a limitations and future directions section in our conclusion.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide derivations of all of our results at the level of rigor of physics calculations.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide example FLAX code which shows how our models are implemented. We also mention the datasets and hyperparameters used.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide code in the uploaded supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We describe the number of ensembles and general experimental details (SGD vs Adam, depth, width, head number etc) in all of our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide errorbars wherever appropriate, usually averaging and measuring standard deviation over different initializations.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide a section mentioning the compute resources used. All experiments can be performed on a single Nvidia A100 or H100 GPU. The details on run times for each experiment are provided in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We are conforming to the code of ethics and preserve our anonymity.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This work is of a theoretical nature and is about general scientific understanding of transformer models.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We are not releasing any data or models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the papers that introduce the common crawl (C4) and the CIFAR-5M datasets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets are introduced in this paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We do not perform any crowdsourcing experiments or research involving human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: There are no study participants in our paper.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper analyzes scaling limits of transformer models w.r.t. key-query dimension $N$, head count $H$ and depth $L$ using dynamical mean-field theory. For the $N\\to\\infty$ limit it is shown that $1/N$ scaling for the pre-attention scores is required for stable learning and all heads become degenerate. Conversely for the $H\\to\\infty$ limit the kernel of each head is shown to follow independent stochastic processes. For $L\\to\\infty$ a branch scaling of $1/L$ is shown to be required for feature evolution. The theoretical results are complemented by experiments on natural language datasets.\n\nSTRENGTHS:\nWhile I am not familiar with the specifics of DMFT, the paper seems to provide a solid analysis of a problem of practical importance - the effective scaling of various hyperparameters in large-scale limits, extending previous muP type works to the transformer architecture. The results theoretically establish appropriate scaling regimes which guarantee diverse kernels are generated across attention heads and are backed by heuristics and detailed experiments on vision and natural language transformers where necessary. The analysis is able to provide concrete descriptions and recommendations of scaling regimes despite the generality of the framework, for example the effects of MLP layers, learning rate scaling and LayerNorm are also accounted for in the appendix.\n\nWEAKNESSES:\nSee Questions.\n\nQUESTIONS:\n* Besides specifying the scaling exponents which seem to be somewhat intuitive from the central limit theorem and law of large number type limits, can the explicit limiting kernel derivations (equations 5-10) be used to predict more detailed aspects of the training process?\n* There should be discussion on the effectiveness of each scaling with respect to compute. For example in the experiments, which configuration was optimal in terms of flops and how well does the theory corroborate this?\n* The distinction between $\\alpha_\\mathcal{A} = 1$ and $\\alpha_\\mathcal{A} = 1/2$ is explained in Appendix C.4, what happens if $\\alpha_\\mathcal{A}$ interpolates between the two extremes?\n* What are some previous results or potential approaches towards analyzing limiting dynamics of neural networks under Adam?\n\nLIMITATIONS:\nA Limitations section is provided in Section 5.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

comments: SUMMARY:
The authors investigate multi-head transformer dynamics by scaling to infinite limits in key/query dimension, heads, and depth respectively using dynamical mean field theory and discover different statistical behaviors.

STRENGTHS:
- Give detailed analysis (and closed form) on dynamics of the updates
- Conduct experiments in realistic settings

WEAKNESSES:
- The paper is hard to understand, the notations are convoluted and most of the community might not be familiar with dynamical mean field theory , the authors might want to offer more background in the main text

QUESTIONS:
1. In line 78, it is mentioned that gamma_0 controls the rate of feature learning, but I couldn't find it in equation (1)
2. Why according to Table 1, SGD's learning rate scales with N and H, but Adam's learning rate scales down when N and H decreases, any intuition to understand this phenomenon?
3. Why 3/2 appears in eq3
4. Why the limit can be straightforwardly computed from the saddle point of DMFT action?

LIMITATIONS:
Limitations are addressed in the last session of the paper
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

comments: SUMMARY:
The authors identify parameterizations that lead to nontrivial feature learning as the limits of $N, H, L \to \infty$. Specifically, the study demonstrates the following:

- Under the limit $N\to\infty$ the $\mu P$ rule is required, causing all heads to collapse into the same dynamics.
- Under the limit $H\to\infty$, each head becomes statistically independent.
- The scaling under the depth limit $L\to\infty$ is also analyzed.

STRENGTHS:
- To approximately preserve the magnitude of the parameter updates, the hyperparameters (e.g., step size) must be carefully chosen based on $N,H$, and $L$. The authors determine appropriate step size scaling for SGD and Adam to ensure feature learning during training.
- The observed degeneration of multiple attention heads under the limit $N\to\infty$ is consistent with empirical observations for $\mu P$.

WEAKNESSES:
- It is unclear whether considering the limit of the number of heads $H$ is practically relevant, as $H$ is typically set to around or fewer than 100 in real-world applications. This raises questions about whether the theory accurately explains the behavior of large language models (LLMs).
- The following paper is also relevant:

  - Lénaïc Chizat, Praneeth Netrapalli. The Feature Speed Formula: A flexible approach to scale hyper-parameters of deep neural networks. 2024.

QUESTIONS:
- Line 78: What is $\gamma_0$? Could you elaborate on this further?
- What is the maximum model size (i.e., number of parameter) used in the experiments?
- Do the results hold for any time t? If so, this feature is worth emphasizing, as most papers, such as [L. Chizat and P. Netrapalli (2024)] and [G. Yang and E.J. Hu (2021)], focus on the initialization phase.

LIMITATIONS:
Limitations are addressed in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study transformer training dynamics under various limits (infinite embedding dimension, infinite number of heads, infinite depth). They point out interesting and subtle behaviours that can happen in these limits. For example:\n- taking the embedding dimension to infinity can make heads redundant with each other, while fixing embedding dimension and scaling head count avoids that issue)\n- 1/depth block multipliers have less interesting kernels at the start of training than 1/sqrt(depth) but this doesn't matter after sufficient training\n\nThe authors also present technical calculations of various kernels that emerge in these limits, derived using path integrals. The authors are honest that these calculations are only expected to be relevant when training time is small compared to the other dimensions in the problem.\n\nI'm going to be honest that I have not attempted to parse the derivations, although I have experience in the area and the results and conclusions all seem plausible and very interesting to me. Even if some of the calculations turn out to be incorrect / make flawed assumptions, I think the paper is still interesting to anyone in this subfield.\n\nSTRENGTHS:\n- the analysis and subtle issues raised about different ways of taking the depth limit (what block multiplier) and different ways of taking the attention limit (heads versus embedding dim) are very interesting\n- the plots in the paper are really interesting, and should be of interest to someone doing practical transformer training\n- the theoretical analysis involving path integrals may be an important contribution, however I am not sure here\n\nWEAKNESSES:\nI am going to provide feedback here that is intended to be constructive. I think that sometimes I can have a blunt style so I just want to start by saying that I think this paper is really cool and it was great to read it. But here goes...\n\n### **Technical tools**\n\nThe authors apply a path integral formalism to derive their results. A major question I have about the work is \"are path integrals necessary here?\"... Why do you use them? It would be worth adding a section to the paper, perhaps to the introduction, explaining the relative merits of the path integral approach to other approaches that people are doing. Ideally this section should be easy to understand by someone who doesn't already know about path integrals!\n\nNow I have physics training myself and found the path integral approach to theoretical physics very beautiful. But there are many ways to skin a cat. For example, I recall that you can take the differential equation $\\ddot{x} = 0$ which is trivial to solve and is the subject of high school calculus, and throw path integrals at it. Is that what you're doing here? Or is there actually a reason why path integrals are the right way to solve this problem?\n\n### **Exposition and clarity**\n\nI think there are a few ways you could improve the exposition. First of all, another round of proof reading. E.g. $\\gamma_0$ at the bottom of page 2 is undefined. Second, I think that whenever you mention a variable, e.g. $\\mathcal{H}$, you should always precede it by the simple English word, e.g. \"the number of heads $\\mathcal{H}$\". This helps the reader parse the paper---remember that they don't have all the same mental variable bindings as you, and it's especially important in a paper like this where there are about 15 different variables floating around. Third, you should make sure that your figure captions can be read and understood by an informed reader in isolation from the rest of the paper---I think again the main problem here is that Greek letters are often used without English signposting. I also want to point out that your plots are missing basic labelling. E.g. on figure 4 b) there is no colour bar or scale.\n\n### **Missing related work on non-asymptotic approaches to scaling**\n\nThe authors cite a lot of related works on asymptotic approaches to scaling and learning rate transfer, for example writing that \"Further, theoretical results about their limits can often be obtained using Tensor Programs [14] or dynamical mean field theory (DMFT) techniques [15, 17].\" However, they do not cite or mention a body of work that works out non-asymptotic analyses of feature learning and scaling. For example, this paper arxiv.org/abs/2002.03432 broaches the topic of learning rate transfer and its relation to architectural dimensions a year before the muP work came out (and at a time when most other theory researchers were working on NTK analyses). There are clear advantages to the non-asymptotic approach in that it is easy to apply to different initialisations and different base optimisers.\n\nQUESTIONS:\nPlease see the Weaknesses section.\n\nLIMITATIONS:\nI think the authors do a good job talking about limitations at the end of the paper. I think it would be helpful to the reader to clarify the mechanism of the limitations. E.g. \"If training time was much larger, then even for very large embedding dimensions, initially negligible stochastic fluctuations between heads could gradually amplify and lead to different large training time behaviour than what we describe here\".\n\nBy the way, I am willing to upgrade my score. But will be keen to see and engage with the opinions of the other reviewers.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}