Infinite Limits of Multi-head Transformer Dynamics
Authors: Blake Bordelon
Hamza Tahir Chaudhry
Cengiz Pehlevan
Venue: NeurIPS 2024 poster (2024)
Keywords: infinite limits
transformers
mean field theory
learning dynamics
Primary area: optimization_for_deep_networks
TL;DR: A theory of various large model size limits for transformers

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
Blake Bordelon, Hamza Chaudhry, Cengiz Pehlevan
John A. Paulson School Engineering and Applied Sciences
Center for Brain Science
Kempner Institute for the Study of Natural and Artificial Intelligence Harvard University
Cambridge, MA 02138
blake_bordelon@g hwrvard. edu
hchaudhry&g. harvard edu
cpehlevan@seas harvard. edu
Abstract
In this work, we analyze various scaling limits of the training dynamies of trans- former models in the feature learning regime. We identify the set of parameteriza- tions 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 which infinite limit is taken and how attention layers are scaled. We provide numerical evidence of convergence the limits and discuss how the parameterization qualitatively influences learned features.
1 Introduction
Increasing the scale of transformer models has continued improve performance of deep learning systems across many settings including computer vision 2 3 and language modeling [5] 6 7
[.. However, understanding the optimization stability and limiting behavior of these models under increases in model scale remains a core challenge.
One approach scaling up systems in stable and predictable way is identify parameterizations of neural networks that give approximately scale-independent feature updates during training [T0] [2] The mean field parameterization, commonly referred as μP, is well-known example that satisfies this property [13 [4] [5). When such parameterizations are adopted, the learned internal representations hidden layers of the network very similar across model scales [16] [7], but performance tends to improve with model scale [10 [2] Further, theoretical results about their limits can often be obtained using Tensor Programs
14 or dynamical mean field theory (DMFT)
techniques 17
In this work, develop theoretical treatment of randomly initialized transformers. We study various scaling limits the training dynamics these models including the infinite key/query dimension limit, the infinite head limit, and the infinite depth limit. Concretely, our contributions the following:
1. We derive DMFT for feature learning in randomly initialized transformers with key/query dimension N, attention head count and depth L. From the derived DMFT action, we identify large N, large H and large L limits of the training dynamics.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

2. We analytically show that the large key-query N - oo limit requires the μP scaling of key/query inner product with 1/N, even if keylqueries are reparameterized decrease the size of their updates from gradient descent.
3. From the limiting equations, we show that this N + oo limit causes multi-heat self attention trained with stochastic gradient descent (SGD) to effectively collapse single- self attention since all heads follow identical dynamics.
4. To overcome this limitation, we analyze the infinite head H. + oo limit while fixing N. We show there limiting distribution of attention variables across heads at each layer throughout training. Despite N being finite, the infinite-head H - oo limit leads to concentration of the network's output logits and learned residual stream feature kernels, giving deterministic training dynamics.
5. Finally, we examine large depth limits of transformers with residual branch scaling. We illustrate and discuss the tension between parameterizing a model that it has non-trivial kernel at initialization while maintaining feature learning within the multi-head self attention (MHSA) and multi -alti-layerrreee perceptron (MLP) blocks.
1.1 Related Works
Hron et al. [8] studied the Neural Network Gaussian Process limit of multi-head self attention in the infinite- head H -> co 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 Gaussian process, establishing a connection to kernel methods. Dinan et al. [19] develop a similar theory of transformers initialization and compute the Neural Tangent Kernel associated with this architecture as the dimensions per head N + 0o using a VN 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 + oo (though the distribution of residual stream variables generally becomes non-Gaussian).
Several works have analyzed the signal propagation properties of transformers initialization large key/query dimension N and large depth L [20 21 (2 (3 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 with a₁ E [ 1], which has been shown to converge to but also throughout training in the feature learning regime 27 However, we argue that in transformers that ag = is preferable as it enables the attention layers to update non nogliggggy as L 00.
limit not only at initialization 25 26
Yang et al. 10] introduced the juP scaling for attention layers which multiplies the key/query inner product with N rather than the more commonly used VN (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 non-asymptotic approach to optimizer design and scaling based on controlling the norm of weight updates (28). This approach coincides with jP width-scaling when the spectral norm of the weights used as the measure of distance [29], and can achieve hyperparameter transfer for 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 dmoste Hdheal H.N which can be increased by scaling the number of heads or the dimension of each head, where N is often written dhead- Our goal to determine the set of parameterizations that allow the attention layers to undergo non-trivial feature learning in the various N, H, I. - ou limits. The analysis of these limits performed with batch size and number of training steps fixed while the other architectural parameters taken to infinity.

---

2.1 Model Scalings
The network's s output is computed by a depth L recursion through hidden layers & E [L] starting with the first layer h (x) VD 2 RRN where a IR is the input spatial/token position s.
Preactivations in subsequent layers he are determined by forward pass through the residual stream which contains an attention layer and a MLP layer
h E+1 h
80 L
MLP (n3) h - h
80 MHSA (h'). L
(1)
The constants "oo and Bo control the rate of feature learning and the effective depth respectively T will consider og E 2: The multi-head self attention layer (MHSA) with pre-layer-nomn
We
MHSA (h),-
Σ bE(H
tor 06
Σ s'E(S]
of o's
On
06
VNH
0/8
W Wi
=
(2)
VNH
where of E RRSS are the attention matrices passed through matrix-valued nonlinearity For a sequence of length S, the pre-attention matrix A
(Ai)
SxS is defined as
1
1 W h t N3-
Abse
kiS NI- VH.
W Q6 h
(3)
Noa
The exponent O,A will alter the scale of the pre-attention variables Ab at initialization. The input matrices have shape W W W Qh
E NRNN while the output matrices have shape W
Kh.
Kb are initialized with ( entries while W Kb W Qb have entries of size @(N1-a4) which ensures that all key and query k q vectors are (1) at initialization. The pre-attention variables A RSXS each head are determined by key kg and query &ns' inner products. The MLP layer consists of two linear layers with an element-wise nonlinearity applied in between, where W W NHXNN NH are initialized with ((1) entries:
NNHN All of the weights Wor W Qu: W
MLP(h.) VNH
W
(
h
1
W
hg h LN (nb))
VNH
μP scaling parameterization
downscales the readout of the last layer compared standard and NTK Thus, we define the output of the model as
32
wL (5 h
f No
(5)
5 where h E N are the final layer preactivations spatial/token position [S]. The parameter To 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, must scale the learning rate y appropriately with (N, H. L). However, this scaling depends on the optimizer. In Table 2 provide the appropriate scaling of learning rates for SGD and Adam for any ar [ and
a E [2 1). In what follows, we focus on the SGD scaling and theoretically analyze the N - oc, H - so, and L > co 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 ensure convergence to the L + co limit. Example FLAX implementations of these parameterizations for vision and language modeling transformers are provided in Appendix
'The scale of the update the residual stream due each layer will be 077088/L)
"If az < 1 or a.4 < some of the forward pass variables would diverge at initialization as L + co or N + oc respectively. If az > 1 then updates to internal residual blocks will diverge as T. + co.
'The LN denotes layer- which we let defined terms each vector's instantaneous mean and variance I.((h) = v
(h- μ1) where μ = Nम and o² = N99 11
"The nonlinearity generally softmax. For autoregressive tasks, is taken to be causal.
Instead of pooling over space, outputs f could also carry spatial index (such next word prediction).
3

---

Headss H. D0
An
Now Ks.
A
Bg
Amention
Valacs
Ouaqut
(a) Scaled Residual Stream
(b) MHSA Block
Figure I: Schematic representations of the transformer architecture we model. (a) The forward pass through the residual stream is alternation of MHSA and MLP blocks scaled by Byl-ax, (b) The MHSA block computes keys, queries, values, and attention variables to produce a concatenated output of dimension deddet NH.
Optimizer
Global LR
First/Last Layer Rescale Multiplier La -i-ai.
First/Last Layer Std. Dev.
SGD
NoN 201-1
9
Adam
To
L. 1+02
T NN
1+0L
VNH
VNH
Table I: The learning rates which should be applied to obtain the correct scale of updates for SGD
Adam optimizers. In addition, the weight variance and multiplier for the first layer may need to be
rescaled (relative eq (5) with width/depth depending the parameterization and optimizer.
Our analysis assumes that at each step / of SGD or Adam a mini-batch B, of size O(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 fixed number of steps or full batch GD (repeating data) with a ((1) sized dataset.
3 Infinite Limits of Learning Dynamics
In this section, first analyze the infinite dimension-per-ie N - limit of training. We find that for this limit, the yP rule of a dynamics. To counteract this effect, we next analyze the infinite head H. 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 + co limits.
is necessary and show that all heads collapse to the same limit of the training
3.1 Mean Field Theory Treatment the Learning Dynamics
To obtain the exact infinite limits of interest when scaling dimension peroo---p-eeaddd N, the number of heads H. or the depth L to infinty, we work with 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 33 35| 36 37| 38 (39. Following [1 L.]. we use this method to reason about the limiting dynamics of randomly initialized neural networks by tracking 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 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
One way of obtaining a well-defined infinite parameter limit of transformers is to take the N + oo 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 श k q. The keys and queries are uncorrelated
initialization which motivated the original choice of a 2 55 [8] Yang et al. [10] assume the entries of the key and query vectors move by (11), implying 0 = 1 is necessary since the
Training time, sample size, sequence length/spatial dimension all treated as fixed @(1) quantities.

---

10'
Test
Vars
10-2
1.0
10-4
14-2
0.8
10t4
10m2
10-1
100
101
10'
10*
Learning
(a) Hyperparameter Transfer for Various 0,A (b) Attention variance across heads
Figure 2: Increasing dimension-pers per-sson---e--aadd N with heads fixed for a (1. (a) Both A
and a = exhibit similar hyperparameter transfer for vision transformers trained on CIFAR-5M over finite N at 7 16. (b) The variance of attention variables across the different heads of vision transformer after training for 2500 steps on CIFAR-5M. For a = the variance of attention the variance does not decay with N.
variables decays rate ((N- and for a 2
update to k is correlated to and vice versa. However, it is possible to obtain ((1) updates to the attention variable for alternative values of a.4 we choose the change key (also query) entries after
gradient descent to be ōk; ~ O(N- ). 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 F.1.2 any well defined N + oo limit of SGD requires a The reason 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 N + oo unless initial key and query weight matrices are downscaled to have variance of order N (1). In Appendix we provide DMFT analysis the + oo limit of the transformer training dynamics. We summarize the result of that analysis informally below.
Result 1 (Infinite Dimension-Per--Hee N) (Informal) A stable feature learning N co limit of transformer SGD training requires taking = 1 (μP scaling), even if key/query updates are allowed be rescaled to account for their correlation. The limiting dynamics of training are governed by the residual stream kernel H (x, x', ,t,,') NH h (2, t)- h,, (t) and a collection of inner product 00
kernels in each head that concentrate as N
1
V Vhoor (x, x', "t,t,t" = "ha(x, t)- (a', (f = N 4ho (ho
(6)
N 1
Kt hee (x, x' ,t,t')=
N ke (2, (a', f), Abse (x, t)= K/ (xx 1) (x.,t.
alongside residual-strean gradient kernels and response functions in the of [15) [7y. The NN output logits f(x, t) evolve deterministically according to the above kernels as well as kernels for the gradient vectors g NooN 8f which appear in the backward pass. These variables become Sh
identical across heads such that for any h, ly' E [H]. Absor (x, t) the residual stream and key/query/value variables within a MHSA block are statistically independent across neurons and can be described by single scalar stochastic process
(2,,), (2, t). All preactivations
1x,, t) (x, + 3oL OLL u (x, t) BoL uu
+ +
Σ Σ
do (z, x', t, (x', t') (2, (t'
8 (
kh (2, ukxaa (2,,
Σ
/ dx' Che (2, x tt tht (x',t)
(8)
where i. ue .u'p are Gaussian processes with covariances Hii1, Bii,VVoo He respectively. Analo- gous equations hold for the queries and values. In the limit, the kernels Hag x', tt,,,,)' (o((x,t)), (),5(a) (x', (x',t') Aluo (x, (Kf, (x, t) (x, t)). etc. computed averages () over
5

---

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 Ct C Oke depend on the precise mini-batches of data B, presented at each step which we assume are known.
We derive this result using a Martin-Siggia- Rose path integral formalism 40] for DMFT in Appendix Full DMFT equations can be found in Appendix E.2 Following prior works on DMFT for infinite width feature learning, the large-N limit can be straightforwardly obtained from a saddle point of the
DMFT action [T5] [i [7].
N oo, multi-head self atlf-attention will effectively compute the same outputs as a single-head self-attention block. We theoretically derive this effect in Appendix E.2.T and demonstrate it empirically in Figure (b). This experiment shows that in 0.4 = 1 (μPP) transformers trained for 2500 steps on CIFAR-5M [2). the variance of attention matrices across heads decreases with N. However, we note that if the scaling exponent is chosen instead as a = there is von-decraas diversity of attention variables across heads. This result is consistent with recent empirical findings that attention variables in uP 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 large-N N 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 there exist well defined limits at finite N with diversity of attention variables across heads.
Collapse of Attention Heads
3.3 Scaling Number of Heads
In this section, we take H -> oc 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. + co limit of SGD training dynamics in ran- domly initialized transformer at any key/query dimension N, scaling exponents A, O,L, and any depth I. is governed by head-averaged kernels for pairs of imput data x. xt at training times L, and spatial/token positions s, such as
vlo (z, (x,x't,,))= NH
1
Σ h=1
hh
Ea (x,
Ea -bs (a',t)
(9)
Co. The attention variables (K (ax t), (xx,)) t), (f A), (x. f)} within each head become statistically independent across heads and decouple in their dynamics (but not across dimensions within head). Each residual stream neuron becomes independent and obeys a single site stochastic process analogous to Result 7 but with different kernels.
which converge to deterministic values
H
showing that the joint distribution of head- averaged dynamical quantities satisfies a large deviation principle and the limit can be derived as a saddle point of a DMFT action.
We derive this and the full DMFT in Appendix E.3
To gain intuition for this result, we first examine variables H² and Af, at initialization. In Figure 3
we plot the convergence of a N = 4, L 8 vision transformer's 's residual stream kernel H to its H + Do limit at rate (HH-1) in square error, consistent with perturbative analysis near the limit 17 Next, we plot the distribution (over heads) of A, at a fixed pair of spatial/token positions for fixed sample. This is a non- Gaussian random variable for finite N, but as N + oc the distribution of A will approach a Gaussian with variance ((N11-2AA) ).
We then investigate training dynamies as we approach the + oo limit. In Figure 4 (a) we show the test loss on CIFAR-5M as 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 - oo limit is essentially uncomputable, we approximate it as the ensemble averaged predictor of the widest possible models, a technique used in prior works [16 IT]. We again see that at early time, the logits of H head models converge to the limit proxy at rate OHH-1), but after continued training the convergence rate weakens. This effect

---

10
10
120
216
Densty
133
2
2
A
(a) Initial H' Kernels,
(b) An Distribution
(c) An. Distribution
Figure 3: The initial kernels converge as H > oo and are determined by (possibly con-Gaussian) distributions of A over heads in each layer. (a) Convergence of H (æz, x') = मा h (a:) h (x')
in L 8, .V vision transformer initialization rate O (b) The density of Af, entries heads fixed spatial location converges as H but is non-Gaussian for small N. (c) As
N + oo the initial density of A approaches Gaussian with variance of order O ((N))aa).
7 16
10
2.0
N 64
1.8
N 128
10"
Loss
256
1.6
51512 512
10"2
steps=25
1.4
2t 2048
stops-100
steps=250
1.2
10-2
stops-500
stops=1000
steps=2000
10-1
500
1000
1500
2000
2500
10a
10
10'
Steps
(a) Training Dynamics Varying
(b) Convergence
-> oo limit
Figure 4: Approaching the large head limit - oo in early portion of SGD dynamics for a vision transformer trained on CIFAR-5M with (L, N) (2, 4) and (7o 3, a = (0.05, 4, }) and losses averaged over 10 random inits (colored error bars are standard deviations). (a) loss and the variability over random initial seed decreases. (b) The mean square difference between output logits for H head models and proxy for the infinite head model on held out batch of test examples. Following prior works, our proxy for the limit is the ensemble averaged outputs of the widest models [6
increases the
has been observed in jP networks in many settings 16 and theoretical model of this was provided recent work which argues 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 a Below we informally describe the main finding which again uses DMFT formalism and is based on analyses in recent works on infinite depth networks from Bordelon al. [1] and Yang et al. [12]
Result 3 (Infinite Depth Limit) (Informal) The training dynamics for H. branch scaling with Oz E [2. 1) is described by differential equation for residual variables h I)
30 with L-
in layer time lim T, for the residual stream
30 8am L" dua(T', 2x.
ha(T, 2),
Σ dd' dx' S dr'C xx,'' x,,t",ttt)) (t,t))g (T',',
5
+ 700 Fet"
(10)
where the Brownian motion term du, (T, x. t) survives in the limit only if of = and has covariance
(du, (T, æ. (o'dugrr'',z x,,x,'') t²)) = 5(ß - d''ddddr' (o x', t,t') t + Va (T, (T, (T, ,ttt) (11)

---

and the deterministic kernel Css (T, 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 T, + co limit unless o = 1. All response finctions are suppressed at I, - oo unless 0
Below we provide a couple of short comments about this result. The proof and full DMFT provided in Appendix E.4
1. initialization t = 0, the only term which contributes to the residual stream layer dynamics is the integrated Brownian motion Jg du(r') which survives at infinite depth for az = 1 For aı = 1 this term disappears in the limit. The structure of C(T) is also modified by additional response functions at a = 2 [M] which we show disappear for az 1.
2.
The weights within residual blocks (including the MHSA block) can be treated completely frozen for ag < 1 in the L + oo limit, which leads to the simplified statistical description of the preactivations in those layers. However, the residual stream variables h(r) do still obtain (1) updates. At aj = 1 the weights in the MHSA blocks evolve by o₂ causing additional feature evolution in the model.
3. A consequence of our large N and large L result is that the N, L + 0o limit with as (xx t) = 0 for all time t. Thus
(the parameterization studied by [2) would lead to Afs the MHSA blocks would only involve average pooling operations over the spatial indices, despite the residual stream kernels H updating from feature learning.
[Wg,000 100 Wa,000
Test Loss
10"
10a
30"
1030
Compute
(a) Key/Query Weight Changes
(b) Depth Scaling for az E 2
Figure 5: Depth scaling in a vision transformer on CIFAR-5M with os. E ( 1}. (a) The key and query weights move by I/VLI. (b) The compute scaling laws with models fixed width N, H and varying depth L. At large L, the a = (dashed) models perform better at fixed compute.
First, note in Figure that the weights within each attention block freeze L - with ar case but move at constant scale for a = 1. As a consequence, the loss at large L. can be lower in the parameterization.
We can see some numerical evidence for the first of these effects Figure (a)-(b) (a)- training at large L is slower than the base model and the initial kernel appears quite different for since the
where initially
L= 4 and = 64. The initial kernel will decrease scale as - do for as. preactivation vectors lose variance 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 az 1 parameterization is the one that allows attention layers actually be learned in the limit L - oo, but that this parameterization leads 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 a regime where there 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
8

---

2.25
Base
7 128
= 128
L=
2,00
N=1
Initial H x')
1.75
N- 128
1.50
1.25
1.00
Final H4C, x)
0.50
0.25
10a
10a Steps
10a
10a
(a) Training Dynamics
(b) Initial and Final Residual Stream Pooled Kernels
Base
7 128
N
128
64
Initial F2
8-
Initzal Densiy
Final Density
0.4
Final Mi..
0.2
0.0
(c) Spatial Kernels Single Sample
(d) Attention Distributions Before and After Training
Figure 6: Initial and final representations are converging as model scale increases after one pass of training on the full CIFAR-SM with SGD+momentum. The base model is (N, H. L) = (16, 16, 4) and (a-A, OL, So %0) (1, 1, 4, 0.1). (a) The test loss dynamics for one pass through CIFAR-5M. The dynamics are very similar across different head-counts H but the early dynamics are changed for large depth I., consistent with our theory. (b) The initial and final feature kernels after spatial pooling
at the last layer of the residual stream. The initial kernel at large L. is quite different for CxA = 1 due to suppression residual stream kernel across pairs of spatial positions for single randomly chosen input sample. (d) The distribution of attention entries across heads at a fixed pair of spatial locations and data point. The initial variance of A decreases for a = 1 but the update is roughly consistent across N. For both initial and final distributions for An are consistent across N.
Brownian motion on the forward pass, which we explain in Section 3.4 (c) The
a = 2
investigate realistic training dynamics of our LLM as we take the N, L,H. - co 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 E [1, z 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 a final loss around 5, with slight benefit to a = 1 For fixed H and L, scaling N provides a similar increase in This is in line with our predictions 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 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 slight difference for large N. Furthermore, they are extremely similar for large N and large H.
performance to scaling heads in when 0 = 1, but a substantial increase when a
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 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 jP multi-head transformers in the limit of N + oo collapses to single-head self-attention. At finite N and infinite heads H. + co

---

22
(b) Kernels: final token across samples
20
N
Steps
(a) Training dynamics of LLM on C4
(c) Kernels: tokens within single sample
Figure 7: Training dynamics and initial/final representations of decoder only language models
trained C4 converge with increasing model scale. The base model has (N, H. L)) (8, 8, 8,44 and (o1,/o, 7o) = (1, 4, 0.25) and a E (1, )- (a) Train loss dynamics after 10000 steps on C4 using Adam optimizer. The dynamics improve consistently when scaling for both values of a,A. with slight benefit to a 2. Scaling N reveals significant advantage to setting a = Scaling L provides little improvement for either parameterization of a,A. (b) Initial and final residual stream kernels for the final token across samples for Base, H. 128, N = 128, and L = 64 models. The first row initialization. The second and third rows are after training with a {1, 2 respectively. (c) Initial and final feature kernels across pairs of tokens for single randomly chosen input sample. Note both types of kernels are identical across a except for slight difference large N.
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 (az 1) causes the initial kernel to lose structure from the initialization as L > co, but allows learning of the self-attention variables, whereas the depth scaling that preserves structure from initialization (az, }) leads to static layers.
Limitations and Future Directions Currently exact theoretical analysis of the limit is focused SGD (and can be easily extended SGD+momentum (15)) while Adam is currently only reasoned with rough scaling arguments rather than an exact theoretical description of the limit. Since Adam most commonly used train transformers, 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 understand learning dynamics in the regime where model size and training times are chosen to balance compute optimal tradeoff (or perhaps even training longer than compute optimal) 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 an important future direction of research into scaling transformer models.
8 39
10

---

Acknowledgements and Disclosure of Funding
BB would like to thank Alex Atanasov, Jacob Zavatone-Vetth Lorenzo Noci, Mufan Bill Li, Boris Hanin, Alex Damian, Eshaan Nichani for inspiring conversations. We would also like thank Alex Atanasov and Jacob Zavatone- Vetho for useful comments on an earlier version of this manuscript. BB supported by a Google PhD fellowship. HC was supported by the GFSD Fellowship. Harvard GSAS Prize Fellowship, and Harvard James Mills Peirce Fellowship. CP was supported by NSF Award DMS2134157 and NSF CAREER Award IIS2239780. CP is further supported by a Sloan Research Fellowship. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. The computations in this paper were run on the FASRC Cannon cluster supported by the FAS Division of Science Research Computing Group at Harvard University.
References
[1] 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.
[2] 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.
[3] 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 66 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. / /929, 2020.
[4] Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab, Matthias Minderer, and Yi Tay. Scenic: jax library for computer vision research and beyond. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pages 21393-21398, 2022.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[6] 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.0836/ 2020.
[7] 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-1900 1901, 2020.
[8] Jordan Hoffmann. Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
[9] 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.00 08774, 2023.
[10] 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:1 17044 2021.
[11] 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 htttps://pppeerreee bet//forum2id-k2leehiee
[12] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Feature learning in infinite depth neural networks. In The Twelfth International Conference Learning Representations, 2023.
11

---

[13] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-fiee field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pages 2388-2466.. PMLR, 2019.
[14] Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pages 1 11727-11737. PMLR, 2021.
[15] 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-3222 2022.
[16] Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks consistent across widths at realistic scales, 2023.
[17] Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctua- tions in mean field neural networks, arXiv preprint arXiv:2304.03400 2023.
[18] 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.
[19] Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization, 2023.
[20] Yihe Dong, Jean- Sann-Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803. PMLR, 2021.
[21] 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-2728 2022.
[22] Bobby He and Thomas Hofmann. Simplifying transformer blocks. arXiv preprint arXiv:2311.011 2023.
[23] Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli. Geometric dynamics of signal propagation predict trainability of transformers, 2024.
[24] Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and-widd limit. Advances Neural Information Processing Systems, 36, 2024.
[25] Soufianc Hayou, On the infinite-depth limit of knite-widthh neural networks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL htttt::://ppereeiieeeee net/
Forum?dd-RbLsy2kka
[26] Nicola Muca Cirone, Maud Lemercier, and Cristopher Salvi. Neural signature kernels as weinite-width-depthhe limits of controlled resnets. arXiv preprint arXiv:2303. 2023.
[27] Lénaïc Chizat and Praneeth Netrapalli. The feature speed formula: a flexible approach to scale hyper-parameters of deep neural networks, 2024. URL atttps://axxvv.orr/abb///// org/abs/2311 18718
[28] 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, 23:21370-21381, 2020.
[29] Greg Yang. James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:22100.7821 2023.
[30] Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue, Automatic gradient descent: Deep learning without hyperparameters. arXiv preprint 27X4v2230..05187, 2023.
[31] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018
12

---

[32] 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.
[33] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[34] Haim Sompolinsky and Annette Zippelius. Dynamic theory of the spin-glass phase. Physical Review Letters, 47(5):359, 1981.
[35] Moritz Helias and David Dahmen. Statistical field theory for neural networks, volume 970. Springer, 2020.
[36] Stefano Sarao Mannelli, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborova. Passed & spurious: Descent algorithms and local minima in spiked matrix-tensor models. In interna- rional conference machine learning, pages 4333 4342. PMLR, 2019.
[37] 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.
[38] Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zde- borova. Rigorous dynamical mean field theory for stochastic gradient descent methods. arXiv preprint arXiv:2210.06551, 2022.
[39] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. dynamical model of neural scaling laws, 2024.
[40] Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A, 8(1):423, 1973.
[41] Blake Bordelon and Cengiz Pehlevan. The influence of learning rule representation dynamics in wide neural networks. arXiv preprint 20Xiv:2210.02157, 2022.
[42] Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. arXiv preprint arXlv:2010.08127, 2020.
[43] 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 unified text-to-text transformer. Journal of machine learning research, 21(140):1-61 2020.
[44] 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.
[45] Etai Littwin and Greg Yang. Adaptive optimization in the oc-width limit. In The Eleventh International Conference on Learning Representations, 2022.
[46] Mufan Li, Mihai Nica, and Dan Roy, The neural covariance sde: Shaped infinite depth-and- width networks at initialization. Advances in Neural Information Processing Systems, 35: 10808, 2022.
13

---

Appendix
A
Additional Figures
N-15
N-
K-
N-
133
N- 125
Lo55
2.90
1.50
3.59
1.25
1.25
2.25
1.00
0.5
100
20'
Steps
Steps
Steps
(a) Vary H
(b) Vary
(c) Vary L
Figure 8: One pass training on CIFAR-5M with vision transformers with the setting Figure
== 16
H 32
H = 64
H = 128
Initial (x
Final H x
16
N=
64
128
Initial H x
Final H x)
L-4
L 8
L 16
L 32
Initial H x')
Final H [x,x]]
Figure 9: Examples initial and learned kernels in final residual stream layer with various extrapola-
tions of base vision transformer model with (H, N, L. (16, 16, 4) trained CIFAR-5M.
14

---

==16
H
= 64
H
Initial H
Final H
N= 16
N = 32
= 64
N = 128
Initial H
Final H
8=8
1=16
Initial H
Final H
Figure 10: Spatial kernels for single test point before and after training across H. L values.
2.2
N=
2.2
N= 8 N
N
2.0
64
N
Loss
N 128
Test Loss
N 128
N - 256
1.8
N 256
1.6
1.4
1.4
0
500
1000
1500
2000
2500
500
1000
1500
2000
2500
Steps
Steps
(a) N Scaling with
(b) N Scaling with
Figure 11: Early training dynamics CIFAR-5M in vision transformer with different dimension- per-head N with heads fixed at H = 4 for a = [)
15

---

F00
6.5
6.2
6.0
6.0
Loss
N 16
1 16
N 64
N
64
410
10"
20
41
10"
7.0
7.0
6.5
0.5
Base
Bose
5.0
N
5.0
N
4.8
4.5
N
125
N 125
10
10
109
10
100
110
10
10
7.0
7.0
5.5
Dase
5.5
Bose
3.
1
5A
1
1 64
1. 64
Compute
Compute
Figure 12: Performance of language models trained on C4 in main text Figure 7(a) as function of compute, estimated as FLOPs examine scaling up N, H, L with either Q 1/2 or a = 1. The Q = models perform better at fixed compute for either N H scaling. Increasing L does significantly increase compute in this regime since the embedding and decoding layers contribute most of the parameters.
Params. The base model has size (N,
H. (N,
L) (8,8,4) and
16

---

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 fixed layemorm operation
from flax import linen as nn import jax.numpa numpy as jnp
class LN_Fixed (an. Module):
eps: jnp float32 1.0e-6
Cnn. compact
def (self, x): 10
12
features shape [-1] number of features
13
mean jup mean x axis -1 mean of var jup var x axis ) var of x out (x mean i newaris] ) jup sqrt( var jnp.
14
15
newaxis] + self eps
16
The MHSA layer is implemented the following where scale_exp represents d'A.
MHSA attention layer
from einops import rearrange
class Attention (nn Module):
"m"Muliihead Self-Attention Layer""*
scale_exp: 3np.float3. dim: int heads: int
10
def setup(self):
self.c 1.5 self scale_exp exponent for the scale factor kif_qk nn. iniiializzie waw.aliiers.noomllce
13
normal (stddev self dself 0.5) ) possible scaling with N
14
kif mn. ww.iiiiaiizzee... normal (stddev = 1.0 ) # 0_N(1) entries
15
computes key, query, value
16
self heads self dim, sernel_init, kif_gk, use_bias False)
self la_eeyer mn. Dense (features
self v_laer nn Dense (features self heads * self dim, kernel_init kif_v, use_bias False)
18
nn Dense (features self heads self sem init kif_v, use_bias False)
self. out
19
return
def call_ (self inputs):
22
23
qk self lk_layer (inputs) / self heads** ** (0.55) self dim+ self. .)
qk rearrange qk, "b 1 (h d) b h d h self heads) (batch, heads, lac,
25
9, jup.split(qk, 2, axis -1) # gives q, k each of shape batch, heads, loc,
self v_laeer (inputs) jup sqrt( inputs shape [-1]
28
v rearrange "b (h d) A self semf.dim** (-so...ssss....lllee..... scale_exp)
d', self heads) imp.einsum ('ijkl ijml >ijkm
29
9, k) # batch x heads loc loc sigma_A softmax A, axis=-1

---

out jnp einsum ('ijkl, ijlm ->imm>> sigma_d, v) (batch, head d)
loc, d)
32
d out self self.out_layeer (out) jnp sqrt ( out shape [-1] return out
out rearrange (out,
33
34
The two layer MLP block is implemented as the following with ø = gelu nonlinearity.
class MLP
MLP_Block Block (nn. Module):
Tuo Layer MLP Block""
features: int
@nn compact
def
N self self.features
kif in.iiiaaiieers normal (stddev 1.0) O_N(1) entries
h mn. Dense (features N, kernel_init kif, use_bias False
(x) / hap.sqrt (1)
10
h nn. (m..eeu()))
h mn. Dense (features N, kernel_inis kif, usa_bias False
(h) / jnp. sqrt jnp.sqrt
12
return
13
We also allow for trainable positional encoding matrix.
class PositionalEncodiu (nn Nodule):
"*" Trainable Positional aacoding"" d_model int
# Hidden dimensionality of the input.
max_len : int # Maximum length of sequence to expect.
scale: jnp. float32 scale parameter for initialization
def setup(self):
# Create matrix of [SeqLen, HiddenDim] repreaenting the positional encoding for max inputs
10
self. pobeendd self, Salf.paramf ( "pos_ ppo__ddned,
noraa nn. initialiizer normal (stddev
self scale)
12
(1, i+self max_len, self.
d_model))
13
def
- __call__ (self, x, train-True); B,T x. shape
15
16
x == x self pos eoseedbbddi [T] / self scale
17
return x
Each residual block is implemented as the following. Below we show the ag = implementation.
Residual Block
class ResidBlock (nn. Module)
dim: int
heads: int
features: int
L: int
scale_exp: jnp float32 1.0
beta: jnp float32 4.0
10
11
Cnn compact
12
def cal.alllll....... (self x
h LN_Fixed () (x)
14
h Attention ( (dim self dim, scale_exp self scale_exp,
heads self heads) h
15
x x self beta / self.L h
self
h (N_Fixed (x)
18

---

h MLP_Block (features = self features) (h)
self beta self. L h
19
return
Our vision transformer model consists of an embedding layer which is applied to each patch, positional encoding layer, L residual layers each containing MHSA and MLP block, a spatial pooling operation, and a readout.
class VIT (nn Module)
"simple VIT model with
dim: int
heads: int
depth: int
patch_sizes int
scale_exp: float32 1.0
adam_scale: int 0.0
11
beta: jap float32 4.0
12
13
@nn compact
def -_call.. (self, x):
d_model self heads self dim
16
L self depth
17
D =
18
# patchify images
20
x rearrange (x, "b p1) (h p2) с (w h) (p1 p2 c) p1 self patfhpatch_size, p2 self patch_size) (batch, loc,
patch_ch_dim_
normal (stddev d_model** (-0. S*self scames (L/self beta) (0.5 (1.0-self adam_scale) ) #
kif_first= an iniiiaialss normal
0_N (1) entries
kif mn ininiiaiizzeee normal stddev 1.0 ) # (_)(1) entries kif an initializers normal (stddev (L/self beta) ** (0.5
(1-self. adam_scale)
25
26
# read-in weights
(L/self. (* ** (1.0-self (idan_scale))ddmmmee *0 self adamesalle Dense (features N, kernel init kif_first, use_bias False) (x) jnp. sqrt D self patch_sizevvy
28
# positional encoding
PositionalEncoding(( (d_moeel model d_model, len (32// self. patch_size) *+2, scale d_ (-0. S*self adam_scale (L/self. x+ (1.0-self. (d.0=selff.ddmmscaaee) (x)
# residual stream with pre-LN
33
for in range deep..)) depth):
34
x ResidBlock (dim self self.dim, heads self holf.headee heads,
scale self scale_exp, features d_model, beta=self beta, L L) (x)
36
# last norm layer
37
x ()
38
# pool over spatial dimension
z.mean (axis
(batch, d_model)
x (-aell..
adam adam_ nn Dense
features 10, use_bias False, kernel_init kif kif_last) (x) adam_scale) # for mean field scaling
(1.00e1 11.0-0.5+self
return a.w.
19

---

For the causal decoder only model, we need to modify the Attention layer and also prevent pooling over spatial indices before the readout.
class Causal_Attentin (nn Module)
scale_exp: jup float32 dim: int
heads: int qk_ln: bool True
def setup(self):
10
self.c = 1.5 folf.scaleeex,, # exponent for the scale factor initializers normalnomalllc (stddev self (seff 0.5) possibly needs be scaled with kif_v mn. m.inntializers. normal (stddev 1.0 ) O_N(1) entries
kif
13
14
# computes key, query, value self self.qk. qk_
15
self heads self dim, kernel_init kif_qk, use_bias False)
Dense (features
16
velf.v_layer = mn. Demae kernel_init kif_v, use_bias False)
= self heads * self dim,
self. out out_layer nn Dense (features self heads self dim, kernel_init kif_v, use_bias False) return
18
19
def - calll (self inputs):
22
qk self lk_eeyer (inputs) / self heads** self dim+s( self. c) # (batch, loc, 3+h*d)
23
qk rearrange qk, "b 1 (h d) >> bhildd h self heads) # (batch, heads, loc, d )
24
q, jap.splslii (gkit 2, axis -1) # gives q, k each shape batch, heads, loc,
25
self vlayer (inputs) / v rearrange "b (h d) -> h d', h self heads)
shape [-1]
28
29
A 1.0/ self dim* wwe.f.scaleeezppppppppcc scale_exp)
jnp. jupsuinsum ('ijkl ijml->
i ijkm' k) batch x heads x loc x loc
exp_A = jnp. einsum ('ijk1, xll>>ijk'', jnp. (A), jnp tril (jnp- ((v. shape[2) s.shape (2]))))
31
phi_A exp_A exp_A. (axis -1) [)),, jnp newaxis]
32
out = jnp einsum ('ijkl, loc, d)
phi_A, v) # (batch, head, ")
34
out rearrange (out,
35
out celf.our layer (out) sqpt out shape [-1]
36
return out
39
class Lr_Transformer (nn. Module);
40
=**A A simple Decoder only transformer
dim: int
43
heads: int depth: int
44
45
scale_exp: jup. float32 adam_scale: int beta: jnp float32
18
VOCAB_SIZE: int
49
50
@nn compact
def --call_ (self, x, train True):
20

---

d_model self heads
self dim
L self depth
54
kif_first mn initializerr normal (stdder d_model**<00,55
self. adam_ 0(1) entries
(L/self. beta) (0.5 (1-self. adam scam_ssaaee )
kifo - www.iaaaleee.eooaaaaeeeeeeeeee normal (stddev 0.0
kif nn iniiializzeeraaa normal (stddev 1.0) 0 entries
57
kif_last mn. iwi.iiniiizzera......... normal (stddev (L/self beta) (0
(1-self adam_ d m_model**<<1 5+self adam adam_scale)
58
# embed the batch sequence integers
60
x (L/self beta) -0.5 (1-self adam scam_scale) d_model (0.5 self adam_
adam_acaaee scale) nn Embed 1OCAB_SIZE, d_model, embedding.init kif_first) (x) # batch x seq len x
PositionalEncoding (d_miinggddmmdaar (d_model d_model, scale d_model (-0. 5*self 5f adam_scalg, (L/self (0..) *(1-seff adam_scale )) (x)
63
for 1 in range (sele.deeehdd)
65
(N)
66
x x self self. beta/L Causal_Attentin (dim self dim, heads = self heads) (h)
scale exp = aelf.scale_exx
h = LN_Fixed (x)
self beta/L
XLP_Block (featurese d_model) (h)
69
70
x LN_Fixed ()
71
x (L/self. beta) (+-05 (1 self adam_ * nn nn. Dense
(features = self 1OCAB_SIZS, use_bias bias = True, kernel_init kifo) (1 5+seffff....... adam_scale) for mean field scaling return x
72
C Simple Heuristic Scaling Analysis
In this section, heuristically work out the simple scaling analysis justify the set of parameteriza- tions and learning rates we consider. More detailed theoretical analysis for the limit of SGD training is provided in Appendix D where
exactly characterize the N - xe, H and L + limits. We consider taking heads H. inner dimension 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 N,, 201-1 For Adam, the learning rate should be scaled N
Optimizer
Bulk Parameters LR
First Layer Rescale Factor
SGD
L + a
Adam
Ma
11/22/2-1121
Ll-aL
Table The learning rates which should applied to obtain the correct scale of updates for SGD or Adam optimizers. In addition, the weight variance and multiplier for the first layer may need to be rescaled with depth depending on the parameterization and optimizer.
C.2 Heuristic Analysis of Feature Changes Under SGD
In this section consider performing single update single example to all weight matrices.
8w
1
+
U
ET
(12)
21

---

where gf+1 E RNNH and of E R have (1) entries. Thus, computing perturbation to the forward pass we find
8h oh +
[+ FT 9
1/
L NH
L1-ar
NH
1 =ōh + L NN
E
o
z b
getl
(13)
The term in the brackets is (1) and we see that the perturbation from each layer contributes O(L-)). As there are L layers, this will give a total change to the final layer 1/4 that O(1).
For the Attention variables, we note that the
dW Kb Li-ai /NH h
SW Qh
L1-oz /NH k
1
(14)
where Tb. bif E RRN are the query and key for head h and h RN is the residual stream preactiva- tion. We can thus compute the changes to the keys and queries due to changes in their associated weights
8k6
he
Li-aL N1-aa g H
N -
H
Li VNH
6gh N aa V
I1-Oz k ET
h²
[1-ai d 1
(15)
where H1 = Nw h² he ~ O(1) Combining these changes, we find the following update to the pre-Attention variables Ab N k 96
H² LI-aLN ki KH + 1 1 2 N2-20 2 Af (HH))
8,4 -1-auN
oh d
= O -I,,))))))
(16)
since b k. Ω - (1). This update to the attention variable due to changes W W will clearly die out as I, - unless oz
C.3 Heuristic Analysis of Feature Changes Under Adam
For Adam, the gradient of each individual parameter entry approximately normalized by scale
45 Thus the learning rate y sets the size of the updates. This is why we scale the learning rate as which gives the same scale updates to the weights as SGD
1 1 NH
sw g ET
ET
(17)
LI-aI
Again computing the correction to the forward pass we find
8h ++1 8h +
1
9
++ tele
CT
Lar.
NH
Ll-ar
/NH
1
=5h +
-
2+
(18)
I
NH
Similarly our update generates the same scale of weight updates to the key and query weight matrices
e Kb
1
h ET
SW
1
k ET
(19)
Li-ar NNH
Qb
L LI-OL
We can therefore follow the identical argument to identify the scale of the change the pre-attention variables 8.4f 1-z)-az),
22

---

C.4 What Counts as Feature Learning for Attention Layers?
Any parameterization with ON [2.1] will cause all updates Ab and entries of h to be (1) across finite N. The entries of q and & only move by (1) if 0 (nP 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, S. etc) rather than preactivation or key/query vector entries themselves. Table summarizes two example values of a.A which are of special interest for their N + oc limits.
Variance of A(O))
Update to A Update to k. q Entries
0 = (nP)
(N
(
(1)
0,4 =
(1))
((1)
(
Table Two interesting choices of scaling for the attention layer exponent 0,4 which give approx- imately constant updates to the attention matrices An- The μP scaling 0 causes the entries of the key/query vector entries to move non-negligibly but causes all heads to be identical (and all A = 0) at initialization. Scaling instead with 0 causes the A variables to be random but still non-negligibly updated under training.
The choice a 2 allows the variance of A to be constant size as function of while also enabling learning of these variables. We verify these scalings in Figure
[=]
200
20200
10a
((kP))
10-1
10"4
(a) Change in & Entries (SGD)
(b) Change in A Entries (SGD)
10
[=1
1.07
200
200
10-3
0,
10
((5K))
10-1
10-2
10:4
10-1
10:5
10a
30a
N
(c) Change in & Entries (Adam)
(d) Change in A Entries (Adam)
Figure 13: The update (a) key kn entries and (b) pre-attention variables An after t steps of gradient descent for scaling exponents a E {1, }. At the first step of SGD, the updates to the keys and attention variables are suppressed due to a lack of correlation between Wo and the gradient 0f After training for multiple steps, this correlation increases and non-negligible updates to the attention variables occur. (c)-(d) The same but for the Adam optimizer with our proposed parameterization.
23

---

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 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 distribution [37 39 or the random initial weights in deep neural network [5] III. In our case, we are interested in the last example, where the feature learning dynamics of a randomly initialized transformer is function of the initial weights in each layer. In what follows, we will give primer on the main objects which typically appear in DMFT analysis (the correlation and response functions) to illustrate the main ideas of the approach.
D.2 Example I: 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 Gaussian symmetric coupling matrix.
In this example we show that the DMFT path integral computing something non-trivial about the kinds of dynamics induced by linear dynamical system with random matrix. In this linear example, the DMFT path integral encodes spectral properties of the random matrix.
Let's consider the simplest possible example: (t,(t) VN
+ Σドn (t) where Wi, = W is
Gaussian symmetric matrix (GOE). This matrix is fixed while the state h(t) E EV evolves. The path integral appraoch would tell you that in the N + oc limit, every neuron has identical statistics given by the stochastic integro- differential equation
0ch(t) u(()+ + L/' dss(t, ult) ((t,
(h((t)
stt,s) (I(E))(S)) R(t,s)=
(20)
(du(s)
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, 8) and linear response function R((t,s).
A generic result of this path integral DMFT picture is
All neurons (all variables h.) decouple statistically. The presence of all other neurons only enters through "macroscopic" quantities Ctt, and R(t, s) known the correlation and response functions. The distribution of these functions over random realizations satisfies large deviations principle where the distribution over C, R has the form p(C, R) NS(C,R) where S is the DMFT action obtained from the path integral method.
2. Extra memory terms like S R(t, s s))(s)) appear which depend on the state at earlier times 8 < The Markovian (deterministic) system for p(h) W) becomes stochastic and non- markovian after marginalizing p h) ddWp(h] p(h
p(h/W/r W )p( W). would argue these memory terms are not obvious apriori but are systematic to compute in this framework.
Since this toy example is línear dynamical system, 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
Trexp (W - s)) = / ddp(d)c<(e-o)
R(t,s) s) N
(21)
where p()) is the eigenvalue density of W. In fact Fourier transform of our DMFT equation recovers the semicircle law p(>) Im/R(iAA)
VT X²]+ for the eigenvalues.
In general, one can think of DMFT as a more powerful version this method that can also handle nonlincarities.
24

---

D.3 Example 2: Deep Linear Network Updates
In this section 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 + oo in residual network, we first thought about simply scaling depth in standard MLP. Below we show how the proliferation of response terms gives different predicted scaling with I. than if we naively disregarded response terms.
Consider a non-residual linear MLP network with juP/mean-field scaling with hidden layers with N + 0o. Train the model for single step of gradient descent with learning rate n on data point (x, y) with |z|² = 1 and y = 1 and output multiplier 1/70- The forward pass variables h (t) and the backward pass variables g'(t) are defined recursively as
h (t)
1 VN
w (t)h'(t
W (O)n' + moo Σ (t.sgt))
(22)
VN
g g((1) 1 (1 g (1)
1
w W'(0) (1) +170 1710 E 8),
(23)
VN
Now, naively, one may think that covariance (t, t'). However, this is incorrect and the DMFT response functions give an additional correction.
W W' (0)' (t) has entries that are independently Gaussian with
VN
The DMFT Equations for this Model Following the approach of 15 DMFT equations for the preactivations h after step of training
11
find the following
(0) uf v²()) g = r0)
he h1) = (1) + (1, + -
(24)
where the random variables (0), (1) and (0) have the following covariance structure
(u'( = (0,00, (1)= = (1,00,,0), (1,0) (I((1) (1,
(r)())((0) (0,0,
(25)
and the feature kernels (t,)')) (t, and response functions (tt,, 1') have the form u²((ttt) (1)()) (ff (f))
(',,,t') = ( (t)(()
(26)
(rf((t)
These recursions can be solved with the initial conditions (0,0) = and Go (0,00) = which implies that H1, (1, 0) = so that
AL( = myo + +1==((100) En7
(27)
Using this equation, we find
(((,1))((((1))()))) +1 1
k2
(28)
1=1
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 neglected the DMFT response functions and approximated the dynamics as
However, if we had
(1,11) H - +1
H = = 1
(29)
The feature variance after step of gradient descent H' (1)2) after 1 step, the final layer
1 1 1
3 2 2 L 1 228 L
DMFT Response Included (Full DMFT)
H
(30)
DMFT Response Neglected
We see that without the response terms get completely different scaling prediction with L!
25

---

1 residual block scaling (az )) the response functions are Scaled Residual Networks For V
still important to accurately characterize the dynamics and contribute (1) corrections to the feature learning dynamics as L - co. However, for the 1/L block multiplier scaling. the response functions do not contribute in the limit. These facts are nor-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 co of training with SGD. The results can be easily extended to SGD with momentum following the methods of [11 [5]
E.1 Deriving the DMFT Action
In this section we will derive the limiting equations of motion for stochastic gradient descent in the H - co limit. We start by defining the loss function which we aim to minimize
L <= f dx p(x) [x)(
(31)
where p(æ) is the data distribution of interest. We note that this can be the population loss or the empirical loss on finite collection of points. We let (ar,
[x.
(f, represent the error signal on datapoint æ. At each step of training & a batch of examples B, { (t), ( is generated and used to estimate gradient for SGD. In what follows, let Eg, represent averages over the minibatch at time t. We emphasize here that the batches B, are assumed to be given or fixed and are not averaged over random draws from p(x), but rather our expectation simply denotes the empirical mean over the minibatch at time
E [s())]
1
Σ DE
((z).
(32)
B
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 MI.P layers (which were analyzed in prior works, Appendix [E] and layernorm (which in the limit it will only apply 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
8o L
h (ax, t) h 1) +
MHSA (h² (h((x, t)),
(33)
where the MHSA layer is
MHSA (h' (x, t)), =
Σ
ob (2, t) (I, t)
be(H)] d'E(s)
of (x, t)
W (c,,f
(34)
VN
u/ (x, = VNH
We (t)(()(((
1
20' (o, neer (x, a (A,(
Naa k B) (2, t)
1
k/s (x),1)
W (1))((20)))))))
N H
Ths (ar, 1)
N W (x,
(35)
To compute the weight dynamics we again introduce the necessary gradient fields which we previously argued have O(1) entries
0f t)
(x,
N
(36)
ah
26

---

We also introduce the following intermediate quantities which are necessary to characterize the backward pass through the attention layer
M hoar (x, t)
NVH 9. g (x,t) (x,
Do hee' (x,
d bao's arr (x. () = 0 (
(37)
We need break up each of the weight matrices into their initial component and their update from SGD
Wo, (1) oh (0) + 300070
BEwBy
E (ax,())) (ar
L JI-ANNH E'E
BonoNo
(b)(1) (4) W (0) + LI-OLVNH
Σ
(ax, 1' Jo Pons (0, (a.
Pat
55'
Bonono
(t) (k) + L LI
E
BewB
Σ
A(x, t' M (z, (' the (x (z,)))
36'3"
Bunu
W (t) W6n (0) L' NNH
BawB
Σ
(ax (x, (i
t'ct
05's"
We can now express the residual stream as h ++1 (xx t) = h,,
(r. 1) + Bo Xoo (x, 1)
+ Yo Σ
Eain A(x' ,t) Σ
g (2,
(38)
where we introduced the fields
H
Xos
1
Σ
Σ
Xops (ar, (2,, (ons(1,1) 1
W (0)ofssa00))
VH h=1
VN
and the kernel
1 H
H
Vta (2,
Σ
uh (x,
(x', (x',t) f/
Σ
Σ of g!!!
V g' ha (x, ,t,
H.N
(39)
We see that, regardless of the choice of azz the update to the residual stream due to feature learning will scale as L/L which is necessary for stable infinite depth limit (IT]. To track the dynamics of the value vectors, we must also track the variables
XUno
W We, (m)))(a
(40)
VNH
We similarly find the following for the key dynamics
k (2x t) Xkka(xx
8ono70
DEwB
Σ
A(x, t')
(x, x' ,t,,tt" t')
1 a1-a2 N1--N
P't
XK Xkoo (x,
We
(0)
(41)
NH
and an analogous update equation holds for the query dynamics q
Lss (2r, (x,)(
7o I T1-aa.Ni NA
Σ E Egnns
Σ
MhM (a6', (')) (ar 11' 1') (2e, , H' (æ, at ,',,,, t',
v'c
s's"
(2n)(101))
W (0)
(42)
NH
27

---

We see that if the residual stream is frozen so that XK XQ are static, the keys and queries will only evolve the N, L - so limits if ar, a 1 as we argued in the main text. From these equations we can deduce the pre-attention values A (æ:, 1).
Next, we examine the dynamics of the M (z, t) variables which defined as
1
hea (x,
NVH 9. f (x, (2.
We can verify that M2 are all (1, throughout training by expanding the dynamics of the attention output og (x. ().
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, can isolate the dependence of the full dynamics all of the initial weights.
MHSA Layer
We will start by differentiating through the MHSA layer
oh (a, (2, 8)
Σ
9
96
(x, t) +
L
Bu
Σ
Ōhe t)
(H),)))( )
g. g. (x,
9. 1+1 (x,t) +
Bo
ΣΣ
(t) (x,
Lai NNH
b
B0
Bo
We 20 T (x, + LOWNH
Σ Wko Ths (xr (x)
(43)
Lo NH.
Bo
9-
(x, 1)
+
LEgs (x, 1) + (x,
(x, 1)
(x,(tt)) t)]
Lai.
To 70
Exere (a', G) (x, x' s,,t) (2)'",('''t)
L
86 no 7O
Dea' ()) K (x, xl,,", + 2aM (x, (,)) (,, (44)
L
where we introduced the variables
da Dohs (x, t)
Σ
der (ax (2
8¹
LMM (x, 1) = Σ M (x, 1) s (x,, 1)k (x, 1)
(o
hs
3's""
LMa (x, t) = Σ
Maus (z, t h (x, t) t) (zx, t)
Ths
and their associated kernels
Go (2,
9ob (x, t) go (x, t)
NH
1=1
KtM (x, x',t,t'''
LMO (xx t) keMo (a', t') ho
NH
1=1
QeMa (z. x (t,t')=
1
eMO Dna
(x, t) M ( (a',tt) t!)
(45)
NH.
28

---

where we introduced the following random fields
(x,
1
Σ
Sigor (x, ho aa0 t)
VH
ha²
(2, =
WVb (0)
gohs (xx, 1)
VN
Eo (2,
1
Σ Egpo (a, t) no'e"" bor
VH
bafar
W (0) gone (x,
VN
E (x,( i)
Σ
Skron (ze, t) t) (xx t) boo" (x, t)
VH
bo"
1
Ekbo" (2,
W Kb (0)
Tha
(46)
VN
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 Kb (0) and W (0) so that VN W (0)g) (ar, t) has ((1) entries. This means that the entries we originally stipulated in order to obtain & and with (11) entries at initialization. Thus, with this required sealing, we must have either a and (11) variance of the weights, which leads to attention variables which are (NN-1/2) initialization or we choose a = 1 and choose kt q to have entries 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 IP sealing 0,4 to study the N - oc limit. We stress that this effect not visible from a simple heuristic analysis of the forward pass variables after an update like we perform in Appendix C
can at most (1) and not ((N1-9A)
Isolating all Dependence on Initial Conditions summarize the previous sections, we begin by collecting all of the stochastic fields which show up in the dynamics and depend 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 necessary characterize the dynamics of the WH--dimensional residual stream
2 (x,
W on (0)mh t), E
W Wob (0)
9. 1+1 (x, g
VN
VNH
1
(x,)
W (O)h t), (2x
Wt. V. (0) (x,
VNH
VN
(x,
1
W Qb (x, t). (2,
1
W
(0) k (x)
VNH
VN
Qh
(x,
We x (0) (x. 1), (x,
W (0)
ah (x, 1)
(47)
VNH
VN
From these variables, we can reconstruct the entire dynamics for the h. g fields. We there- fore study the moment generating functional for the above primitive random variables. Let
29

---

0, - (Wo,, WV, (0), W Kn (0). W (0))
2H(H = Ea, exp
Σ Ebst
da
Ina (2r, x0,o(1,1) () + f6 (0,,
(ær, 1) Sono(or, (2r, 11]
x exp
Σ
dae
Ek Ina (xe,) xkns( () the (z. 1) Skns(ao,))
Ehst
exp
Σ
da
54 Jas (a.t) Xoso t) + Ins (xx (ana(
Fhst
exp
Σ #hat
dx
The (x, xibolax,t) + Ihe (ax, EVne(e,)))
(48)
We multiply by the identity enforce the definition of these random variables in terms of the initial weights. As an example, we would have for the first random variable X. (x. t) and its corresponding pair (x, t)
Attention Output Matrices
In this section we integrate over W Oh W W
Kh Q6-
i
f dæ TrW (0) Loro (z,,ff + Σ 15
1
goth (z, t) (ax .t) (
In W Oo (0)
exp
VH
Σ
dæda' Xons (2,, t) Xons (z',
t') (t',t')V'o (x, tt,''
2
1t'n6'
Σ
daddx' (x, Eoner (a', (aaa,t/)CG G.G (x, x' t,t')
w/'6
-i Σ
dxda'
Le whs (x, 1) at 16 (x, a 1, t,t))
(49)
tt'
where we introduced the response function
R. gg x6 (x, xd, ,'',,()==
N gfth - f) (onaaa' ,t).
(50)
Value Matrices
Next, we average over the value matrices W (0) which gives
i
E/ f dz TrW Xt
Jônd (xx (eoo)))t
In E (
(0) exp
VN
/ dxedae' (xVns(a,, () XV (x, t) (x,
- '6''s'
Σ
(x, x', 1,(') + (x, El (x', G'nGo xe' ,,,))
- Σ
dxda'' (Xvoll,, (a', R) (x. x., t.,
tt's0'
ph .V (xx ,'',t''') has'
i NVH (2, t) (ar
(51)
30

---

Key/Query Matrices resulting in
Next, we need to perform the averages involving WKO (0) which has entries
In Ww (0) exp (-Vx
i
Σ dx TrW Kb (0)
[ Xkhs ()) + (ho () ((soss
(VH
Σ
doodæ' (Kkoo(t, t) (x 1) te (a, tt,t, it,') + Ekbo (x, t) Exoolaa' (2,, 2t,))
2 tt'
i Σ
dodax'
(xno(o, t) (a', ,'d,))) d') "e,,z,t,, t,t)
t''o'
& &K R
(x,
NVH h (x, (x',
(52)
We follow an identical procedure for the query matrices (O)..
i
/ dx 2o Σ to
VH (a, (2, + (a 1) (xx, 177)
In Bw6 (0) exp
N
Σ
dxdæ
(2bb( Xosa (x. t) H (z, x t,t' (x, t) (abor t') (2,,
Σ
dadda'
Xens(or,t) 1) khe (x', (') R (ax, a' 1,1))
bet'on'
i NVH h2( t) Eg (a',
R
(xc, ,,,,,,''=
(53)
Enforce Kernel Definitions After this step, we can introduce new resolutions of the identity for each the kernels that appeared in the above computation
1 dH (x, xt,t, t'tt))dfi (x, x', t,t')
2wriN -----1
exp (H (x, x t,t') [NHH x',t, h t) (x',
doe as' (n, 1,,, (ar "'',, 1,1)
1
2xiN 71-1
exp (Geo (x. (l,() [NHG (a, a', g () (a', ,1)1
(54)
31

---

This is repeated for all of the response functions which involve sums over NH variables
dl No (x, x', t,,',) (. Eo (x', xx tt,
2wiN-1
Ryyiob (a', a.t ,',)) NRR 16 (', + i
VH g (x,,t (ono
exp
R. hoo'
(x, x' 1, di)dRoo ho's (a', x1,1,
2xiN-1
-Roox (a', a. 1, /NR af, 1,1' + i (2, ( (z',
exp
VH
dR 6
noo (xx, w', 1, Laaa Xk (ad', ax,,,,, x. 1,1)
2viN-1
- Baxxx hs's
(x', x,,',t)
[ xt,,)
VH h5(2,) (a)))))
exp
dR EQ (x, x' t, dd) No
nss
ns's
xt,,t,
2TiN
RR (x' x., t',t)
N ,tt)
i hz(x,t) (a', (1))
exp
VH
There are other kernels which are only relevant within a single head including (0r, Kb. A/: Me). These
dhhse (x, a', ,,,,t"" hes (I,, x', ,,,,,K)
2wiN-1
exp (0, "e,,"'t,t') t,,') [NQ (x, xd xt,,tt) (oo (ho (t]]
dke hoo' (x,, ,,,,, (,, ddd baod (z,, ,,,,,1)
2wiN-1
exp (K hao (z, [NK (x, xt,tt,, - ( (ag E')
dMM hao (ar, 1)
b0o (2),
1=
2viN (2 1)
VH F' () (0 (1)001)
exp
de hoo (Ie, t)dAA (x.
2wiNN1
exp (2, [N - k (x, t) ahe (x, t1]
We now combine all of the order parameters into a large collection Q which are vectorized over all layer, time, spatial and sample indices
Q Vee(H', Ne VV KP,KK'
G V 24
U U T R xa R xx3 UR &t Rh& Rh Lb RM 2t..
(55)
After introducing this collection of order parameters, original MGF satisfies large deviation principle with
Z x fae exp (NHL S(Q))
(56)
32

---

where the DMFT action S(Q has the form
1
danda' (H (xx, xd 1, (')Fte (xc, x ,1,1) + Gar n' 1, (.) (xc xt, ,,)),
st'ss'
H
/ dwdd'' [Chssr (xc x', (') (x,, a (,,') + Kheor (x,, x', ,,,,,,, K)) (x,, 20,,",1,11
HT.
b=1 Fm W'' H
HL ΣΣΣ
dxda'' [Vises (xx x', t, (t,,t)) Vhso (2, x',,t,),,))
H HL
dxc (2, (xx + Mbesr (zx t)M baor (2, t))
Fm too'
H HL
dxdxx [Ry5 (z', x., t, t)ng 26
hoo' (x, (x, t',t') + Ro (x', x,t, tysk (x, t,
h= 8=1 ttlan'
H
Σ
danda'
nepii (a', xx, 1) Eh (x,, (x, ,o,,,',, 1,(') + Rpo's (ad' x., 1', 1
k xx
N'st (x,, a 1,11))
HL
b=1 1=1 1t'ss"
H L
1 In
Z +
In z MHSA.D
(57)
T.
LH
b=1t=1 t=1
33

---

The residual stream single site moment generating functional has the form
d'coro ()dxobo 1) (ar, t)dgvbo () dÉkba( 1)d5kbo (2r, dEcbo AddEqno(a 1)
Zres
II Zhear
2x
2
2x
2%.
dædaa H (z, (z, x', ata,,,,.)))) t,,t)) (x',t)
exp
E= tt'se"
dxod'' Go (x, x ,',,t')) d') (2,))))sss () (z',t)
exp
/
dx dx hear (x, t)ge ( (xx,
exp
266'
E
dxdx' Rors (a. t)X (x', ('' (2, o',,ttt
exp
tt's"'
dxdx' Redxx'Ebbal EV (x, t) t) (a' (oo
exp
tt's!!
Σ Eh tt's"'
daddx' Ekns (x,, t) (a', (a', ()Q (x, x' 1,
exp
Σ Eh et'ss'
danda' (2, (2,, 1) (') (x, ,,,t)
exp
Kor:o (xxons( (2,,1) + Evoo (20,)
exp
16 to
dx (x, (x,, + (x (x,
exp
to
Σ
dxda H &6 (x, x' txt') t)g (al,t)
exp
H Eh
tt'as'
Σ Σ
dada' Rgox (a, al ,1, 1, 1') 1))
exp
H
Eh st'ss'
ΣΣ
dandæ' reg (xx x. (1,"") 12,1))
exp
VH
Et'ss'
exp
Σ
dxcdx R (x, x' ,,,' (x,, ( (z')
(58)
VH Eh
15'66'
Though this expression is cumbersome, we will show that, since A. & vanish at their saddle point this MGF merely encodes the following statistical description of the fields of interest such
(0, ( 1
VH. Σ
dæ' Red S (ax, x', , t.t')g ( (a',t')
v0b0 ~ GP (0, Viser (x, x,,t," tt,t))
(59)
34

---

We note that h only depends on Xo
VR Σ of so it has the form
1
H
Xo
oo of + R. 5o R'
Σ H h=1
1
Rg
o
(60)
H
This fact will be important when we take the large head limit with N fixed in Appendix E.3 Next, analyze the single site MGF for the hidden MHSA layers
d (x, t)dE (x, t) dxvas (z, t)dx dicchs (x, t)dx (x, t) di (a, t)dx (2,
z
II
MHSA,
2x
2m
2m
2m
Ebsur
Σ
dxda' (as
(x, x', tt, t shs (x, t) (x',
exp
te'ag'
dxdxx Kho (x,, æ', t,t') oo (x. t) (ze .t)
exp
to'
Σ
doddx Vns (xx, x' ,, y' (zc, (a
exp
tt'ss'
Σ
dx dar (x,, t)ki (2,, t) (2e,
exp
tool
2 Σ dædæ' (on (), (a', (d) (x, tt,,''
exp
Σ
dæda'' Xi () 1) (x, x' ,,,1))
exp
2
tl'as'
2 Σ
dæda' Rk (x, t) (2 (x, x' ,t,')
exp
tt'
/ dædx' X&so (x. t) (x, H) (a, x' t,t''
2 Σ
exp
11! oo!
i dæ (x, 1) (x, + Xrra (x, (x,
exp
to / dac X2bs(m, (2c, D)xos 1) + Rkos (x. 1) D)k (2,
exp
dxda Ross e+ X6 (x. x', 5% S (x, t)u (x' (x',t)
cxp
tt' "6'
Σ
,t, t, t. 1o.,tt9onn Do (a', (x, x', danda' phi
exp
s6'
Σ
t' 1e' & (x, x', t, Dago (w, t)k (x', dxe Rgg
exp
bot'
daddx' B S' (x, x', (') ()g (x',t)
exp
(61)
hst'ss"
35

---

E.2 Infinite N (Key/Query dimension) Limit
First, we take the N + oo 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 [13 where all of the order parameters Q take on their values at the saddle point Q
8S(Q) @Q Q- 0
(62)
Saddle Point Values for Order Parameters Under this saddle point, all of the order parameters presented will concentrate and 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 and have the form
He (x, ttt' = (1i(2,,))) (x), (x, = (96(2, t) t²)) ,,,), (uhs (xx, (zk,, K (a, x', (') = (k 1) (a),t))
Mhea (x, (x, t)ofs (2x 1))
Dh (xe, (,, (4ho (ax, 1)ubor (a', ('))
Aboo (x, = (1) (I,+) t) (cc,
5g (x,,
RigN (ar 1,() VHH
Suloa (xe'
(00 1)
R (z, (x, ,t,t) VH.
8r br (a',1)
RhSa& & (x, x, t,t' VH.
8ne (x, ()
er Kho (a0
& (x, de, ,,,,,)) VH
Ro5s
br (x',,")
ōu (z,
Ee (x. x ,1,1,)))=
5r br (x',
8gons(x. (xx
noox (x, (,1)
oulvbo (x', (')
5k (x,
H' Xa (xc x11,1')
8 oufono (xe', (a'' 1')
Jaf (xr, 1)
Rd (x. ,ttt' =
(63)
ou (x', .!)
36

---

Single Site Stochastic Processes
Our fields of interest will obey the stochastic dynamics
/ da' Reis (x, ,t,,,, ())g)) (a', 1')
Xe,s(ee (2, () - (x,
VH D'''
(x, GP (0, Vo (xr, x', ,t,t))
c (2, 1) = (x, ()
da' Reee (i, (x(,1)
VH
(2, t) GP (0. Gonao (ax, x', ,t,t
(a,s 1) (x, +
1
Σ
dae' n re xee (xx, (,,, (,d
VH
(x, t) ~ gP (0, 2hso (ior(ar (x,, x',t,' ,t,tt)))
da' rh xee ax,,"' (a',
(x, 1) (as 1) + VH 1t
r (x, (z, t) ~ gP (0, D (I,, x te),tt))
(x, t) = róns( (x, t) + >
/
da' Rass' x0 (x, x', ,t, (e {)
1'''
(2, GP (0, Get (x,,,',,,1)))
(x, 1) = u (x, 1) E/
f
dol Rhee (,, (t,t'))'so (x',
P'a
" ~ gP (0, H (a, ,t,1))
(x, 1) ukb (x, t) + E doe' p (xe, x t, t (xl t)
(x, 1) ~ GP (0, H (2x, 2,,,",,,,,')
1'56 dx' Rgoo (x, x', t,t))a (o (x', (x, t) = u (x, t) Σ f
(x, f) ~ gP (0, H (x, x ,t,,'),,t)))
(64)
Residual Stream The forward pass residual variables obey the following stochastic process
h5t1(2 1) = 11 (ax, () 3o of
(1,
+ 110 vo.00 L- Σ Bx'n A(x' 1) E (d)
(65)
P't
where
(x, VH.
heer t)xo (x. t)
(66)
E'E(S)
the backward pass satisfies
Bo
(x, t) = (x,,) +
(x, t) 28 (x, t) + (x, t)]
Lar.
B6 100000 70
Eag (xr Go (x, a' ,,,,,,, 1') (a) (a', ,t)
L
83 13no% %
Exery ((x', ,t) {)) [K (æ, a,,t,tt,tt, t,t') + Q LMO (ø, t,t))] h6 (x',t')
L
(67)
37

---

The keys and queries have dynamics
kho (x, = (x. t)
+
Be Σ
2" t)) Ho (0, ,t,,)
Pice
6'o"
(2, xina(z,))
Bo
+ -i-ar
Σ B, Σ
A('', hho (x', (a', t'' H (2, x' t,t'
5'="
(68)
E.2.1 Multi-head Attention is Single-Head Attention N - 00
In this section. we use the derived saddle point equations for the N + oo limit and argue that they imply that all heads the MHSA layer learn identical attention matrices and contribute the same feature updates to the residual stream. To do so, we proceed in three step inductive argument.
1. First, we show that initialization, all key, query, value and attention matrices (Qn, Kn, An M2} are equal across heads.
2. Next, we show inductively that these quantities (Q., KG, V.. Ab. M₁ } are identical across heads up to some time, then that implies that the response functions RG are also identical across heads up to that time.
3. Lastly, show that if the response functions Rg identical up to some time, then that implies that the MHSA kernels (Qn, Kn, An, M2} 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
Vh E [H] 2ho (2x, x', 0, 0) Kh (ar x' 0,,0) Vn (xx x', 0) H (a, x' 0,0) Maas' (x. 0) = Abso (x. 0) = 0.
(69)
Next, need 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 XOè with the feature learning update to the residual stream, writing the following compressed equation for the forward and backward passes
1 1) = (x, t)
wh (ax, (x. t)o (x. t)
H
+10008811
dal Chés (x, x² -
(x,
g((x,)) t) = t) +
VH Σ
Yh (x, t)o (x,
har
rener (x,, t) ho (x, t) h ho (xx
VH
ha's"
"knee ( (x. bas" " (z, M Mae (x, t)
H
he"g!"
+ 1o
ΣΣ
dx' dx'Cg' s0/ (I, x², t') (x',t"
38

---

In this section, we compute the infinite head limit with N, L fixed. This limit more technically involved than the N + oc limit which required only simple saddle point of the full DMFT action over all kernels. finite N we cannot use this technique since the kemnels 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 where OK and C9 only involve deterministic head-averaged kernels and thus do not carry index. We now derive useful response function identities
VH Oh t) ouors (x',t")
((E - E - t')0(0 x (x,
dae" Chen ( (x, ,1,"" (
(VH kt (a",t")
+ 1000 k=l 8"
ΣΣΣ
duone (zt'
VH "o"" 7o L E "E ou (x',
L k=t ce g
I.
/
da"Ca" aatr (ar (xx ,t,t")
(vTl
(VH ote (a"" )t") ou (at', ,1)
(70)
These equations give the needed response function R X6. From these equations we immediately see that if hao (xx, 1) = y (xr, 1) (which holds under our inductive hypothesis) then ng H. "x0 This same argument repeated for all other response functions that are computed as derivatives of residual stream variables. We have thus found that
hb. [H]. Ro Xo R xx 5 = =Ry Rh
2h R .. Rh R
Ry
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, = Ab,, M2 My... for times less than f. 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
Aho (x. t) = b (x,t)+ +
ΣΣ
/
dl' Clu (z, . (')g (x',t)
(71)
PiC o'
ΣΣ / da' da''og (z, x', ,,, (') (x',
hn (x, t) = u6pss (x,
(72)
where Cer (ax, x', 1,1") and Cg (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 &y, = Vec(ke
(h (ax, 1) amt and analogously C*K = Mat(C (x. x', 1, (') as in the calculations of Bordelon and Pehlevan [5 41. Using this shorthand, we can express the key/query and attention kernels as
I CK ca
"K + C* u
lh = f6=1-
C*
(w2o + Cd Co'uk))
K I a o* coO
HH + nke H T|
I Cike cM
0 = I C*' C* -1
H + cr H' ]
I od. ckcc ckk -1T
4 I - atK cTT
C² HH'CO'T [I- I-C Caa cke [I- IT
(73)
We thus see that the final kernels K₁ Of Af are all identical across heads. An identical argument can be carried out for the value kernel VV and the Ma order parameter.
E.3 Infinite H Limit
39

---

the previous section will still concentrate as H. - oo including
(2, 1,8)
NH h (2, n (xe
(2,,''t, NH (x, t) (x,
Vel (', H
Σ Vhe ho"""" (x, x', it, t) (x, t) (x' (x',t') 5"
and many correlation and response functions. We will call the full collection of all the necessary averaged order parameters Qglotal. Further, not all of the stochastic fields will be relevant to characterize the residual stream. Specifically, only head-averaged fields x6- 5V Ex. 22 are relevant. For example, the first of these is defined as
SQ
(2, = VH a
Tssr (a,, (xxoss
Soo (z,
(oos)(0, +
1
Σ/ da' Ry' (2, ,z)) (x',
V
VH t's
c'E(S)
-ios (x. t) +
Σ
I do'' Rur (x, x', (
(74)
t'''
We note that the above equation true at any value of but the covariance of uf and the response functions Ry R
random variables However, the residual stream only depends upon collective, head-averaged variables
(x, VH b=1s'E(s]
(75)
E00 (x, x', (,,) = H 6=1 s"
Σ
hea (x, E)
Ro & (2, x', t,t)).
(76)
The intuition behind the large limit that, even though of and Rg f are random variables, there should be central limit theorem for Th and law of large numbers for R 600 (x. x', x'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 <.. . 5g- To characterize this joint distribution, we must also of course keep track of the random fields within the MHSA blocks such as (X2n X X S hE(H]] Repeating the path integral setup of the previous section, we need to performing averages over all initial weights such
In E we (0)) exp
i
Tr Ob (0) (O)T))
Σ dæ (x, (2, (2) (TT + 9, 2 E) ( 17] to
VNH.
dxdxx Lo.(x, (x, t) (o, V')'e' (2,
2
taa
/ dada' (ons(a, Long (æ', ,ild',) 1') (a, 11,,'))
2 Σ
Σ
bE(H) N''ss'
Σ
I doodae (Kool0,)) 2 (a') (t))
Σ
v (x,,
Eo (ad',
NH te'
bE(H)
(77)
40

---

It is clear that from this integral that the relevant self self-averaging order parameters are
1
V (0,
Por hhs
(x, () .0g -gg-(('' (x',t') (')
NH
H
(2, (3, ,,,,,) NH
(x, 1) (2')
h
R (0, x' ,1,1²)
NH 9s ++1 (x, Xos (a', ,')
Re (x,2', NH
Σ 0 (xx t) (a' (x',t)
(78)
We repeat this same procedure for all collections of weights and arrive at the following set of order parameters
Q gloobal Vee(H', G., v K QeMio G Utft', Gt, Dia RIMa Quaa Go) U U 20 Rh Ri'Ek E UR Robxt, Ro'ik, RuSsas
(79)
We expect that these order parameters will be self-averaging since they involve averages over H variables. However, the other variables we introduced (Ab Qn Kb. V₁} will not concentrate at finite N and will instead behave random variables.
After introducing these order parameters we find that the moment generating functional has the form
Z /
dQ global exp (NHL ((Qgloba)
(80)
where S has the form
L
dxdax' H (æ, (t,t))) (z, x',,),,,t) + x', (x, ,t,')]
E=1 tt'
dadd (Vlo (x, x', (, Vit (a, x' 1,1') KLMs (ar x' t,t)) (.tt)R R.Mo (n, (,,))
tt'ss'
dacdæ' darda'/Glo (xx, x', ,,,,,,,,. (,1')Gooss (x, x ,,,'') + 2 (x, ,,,,,,,,,,, 1,11) (ar, x', (,,(1))
F=ll Wo'ss
dxdx [R (z. x' t,t') Ry xx,,t R (z, t, R (æ' xt,
Fml telas!
f dxdx [reli (x, x' (g X' (a', x, f, + R (a, x',,,, t, D xxa (a', x f(,))
tt'
+
In Zres
In Z
(81)
L
NL
1=1
41

---

The single site moment generating function for the residual stream has the form
(xx 1)dR6s (x, () (2, t)d& t) dek, (x,, f)dSks
(xe (xe () dEva(a, (ae, O)dEvs (ax, 1)
Zres=
2
2π
2m
2T
Escot
dadda' [16(0,, (a', ( (')lfy (x, a' ,4,) - (x ((x,t)ss (ad,, (T)6 (2x, a' 6.11))
exp
Ett'as'
Σ
dxdda (Xos(a, (oolart))xowl (ax V'v (z. o' (f,"),,,))) + Ev (xx 1) (z', G' o' 1, (i
exp
ttt'
dada' (Eks(0, (d,))
(0) (2, o' ,111 + E/ (x. 5 10) (")33' (a,,"'
exp
Ettr
cxp
i
Σ
dxdxx'
(os(x,t)g3 2 (a', R (z, x. ,,t)) + C t) (a', ( F (xx 2,,,,
Ett' se'
exp
i Σ
dxdxx
()0 (af, Fi (x. 1,,)) 1t) + (Q (a (1) R (x. x' (..
ttt'
dx
(2,xtt))xo t) + Ev. (x. t) (z)
Nto
exp
Σ fts
dx
Ee
(x, ()o t) + (2 (x, ti]
(82)
We can express the MHSA single-head partition functiion ZMHSA in terms of the remaining order parameters within each head that will no longer concentrate finite N
K MAHSA M ,', V C. A M 0', K2 V.,G6)
(83)
After introducing these order parameters, have
ΣMHSA
dQ'AHSA exp (N (Q
SMHSA
Σ
Bedda''Q's. (a, a. Q' (xx, a', 1,(') + k& (x, 2l,, 1,, Kik aa (z,, x² (1,()))
8'ss'
Σ
dandæ' d (a, x L. (x, 1,,1', Co (æ, x', 1.)') (z, x' (,(')
tt'ss'
Σ
dac[Af (x., t) (zz, t) + M (x. t) M nar (z, t)] + In 2k0
(84)
Iee'
42

---

where the key-qqury- single site partition function has the form deos (.. t)dCos( (z, t) dxv, (x. t)dx t) dxks (xx, t) dxos (a, 1)dx(a(
/
II
gko
2T
2m
2T
2n
5azt
Σ
dxddæ' Qes (a,, x', (a) ()
exp
so'
Σ
dadda' R (a., of (') (a',
exp
ss'
Σ
dxda' Veg (a, (0, o', 1, (o 1) (al,
exp
ss'
Σ
dxddæ (a (x, x', t, uu (xx thug (x', (t)
exp
tt'
Σ
daddæ' Go (ar, x' t, (, (0)900 (v. D) (a', (ee
exp
ss'
Σ
dadd'' K (xr, xd ,,, 1,,')) Mo (ar 1) (x',
exp
ss!
Σ
dxdæ Q2M (xx o' ,t,tt)t) (.)g)) M0 (x, (
exp
Σ
dx Ax (ax th t)gg (I.
exp
tss!
Σ
dada (0, E) (a', ( (a, a',,,,))))
exp
'l'as'
Σ
dæda' XVs (x, t) t,tt)
exp
2
tt'
Σ
dæda' tx)k xz, ,t,')
exp
2
W''s
Σ
dwda'' Kigs (x, ()) (2, (e,,t))le () (x, (x,t,,t t,t)
exp
W''s
dæ (x, t) (x. + X (xx)))
1s
Σ
I
do X2ss (2x, 1)x (2, 1) + 2kol0, ()xks (20
exp
to
dxda! Re"""" xb (x,, x', 1,t1) t) (o), (o
exp
i
1t'ss'
i Σ
dxdax R (z, a',
(z',
ao'
dxdxx R A /
i Σ
Ros" (a, x', ,,',tt') Ret') t)ke (x', (
hat'
exp
i Σ
daddæ' h (x, x. ,,,,,,0 1,1 RRkoo (xc, 1) (2' (1)
(85)
het'ss'
43

---

The saddle point equations for this limit are computed as derivatives with respect to Qgiobal only. reflecting that avad-averages will converge as H + Dc
as 0.
(86)
giobal
The final saddle point equations are given in terms of averages over the distribution of heads defined by ZMHSA which we denote as MIISA 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) (( x', (,,') (e,, (x', (')) (',,,'', (o (ze, 1)gg (0', ,1)
V ( x') = (V' (x, a., 1,1)) (a, t)o ()),
C (ar, x²) = (G ,,,,,,,,1,, 1o (ar, (o, (a', (x²,())) MHSA
K2M (z, x¹) = (K'( d' t, of (')M' (x, t) (a', (x) (x, ()' (x',t)) (E)) MHSA
(2N x') (Q' a' L, ('M (x0, (a', (t)"""")'')))ee))) (x, (a' (0) MHSA
R 2b (I. (1,1'
8g. (ar t)
ou (x''
E (a. al, ,1,1)
tt,
(
(0g(,r, ()
RA 3o'
(ax, x ,t,))
orko (x',1))
s (x, a ,t,t'
8h1 (ar, 1)
39'
8r (x',1")
Re oss
So (xx t)
(z, t,t)
Tr
N
dr dr (xt' 1T)
MHSA
(2,
Reoxi (0, ,1(' N Tr
Suivno (2),t))
MUSA
Xaaxa (x, x ,,t, = 1 Tr
1
Sk-Ma (xx, oulg (I),t))
N
MHSA
Re (a, x. 1,11) = N
1 Tr
oghs (z),)
(87)
MIISA
Residual Stream Dynamics
The residual stream satisfies the following single-site dynamics
(æ, 1) + 1,51(e (x. () -15(2x
(x.
30 dæl Re (a, ,,,,,,)))) ( (a', (!)
Σ
L 835
Yoro L
Σ
Egr Vig (xx x', (1'))==+ (a', (t)
~ GP Von)
v<l
Bu
((xx t) =g5 -/-11 (x, t)
Los (Z,)(z))) + (x, t) (z,
80
dae' Relox (x, 1,t') + R x4 (a, x' (,,'') R x (a, x', 1,1)) (a',t)
+
no 20
Exr n3 (G (z, 1,, + K (z, x' (() + MM (a, a' ,,, (a)
P't
(x, 1) ~ (P ()) (2, QeM6), (m, GP(O, Q
(88)
44

---

This matches the result provided in the main text which introduces compressed
Cza x ,'t,t')
R
(x, + pt A'())((' (a(z't))Va V' (a, 2. ,t,t)
(89)
7o
where pr(x:)
1B Exr 6(x x') denotes the uniform distribution over the batch B,.
E.4 Infinite L Limits
In this section, we discuss the two large limits. This can be derived formally two distinct ways. First, one could start with the imitial For this section, it suffices to reason about the scale of the Gaussian noise which appears 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. I, - do limit, we use the fact that the random variables x" u² + R$ g' decompose into a Gaussian ue which uncorrelated across layers and a linear response R² are response functions. This implies that
h L
30
30
R
32 TJO 70,20
Σ
ko g
(90)
L
La
L
1==1
k=1
k=1
We first note that the sum of the Gaussians is a zero-mean random variable with standard deviation
1
Outrow
(91)
Thus, this integrated random variable will vanish unless 02 = 2-
Next, we can investigate the scale of the residual stream response functions. For instance
th
dh Or = OL
dg Duk O (L (L-ar)
0g
O (L (L-a))
0 (L aL)
Fu
Dr
R = 0
(92)
Lal
As consequence, we see that the effect of the Gaussian and linear response terms will vanish as L + provided that az > 1.2 We will consider first, the case where which gives an ODE like limit for the residual updates before moving onto the more involved a Y/ case.
To formally take the - oo limit, we redefine all of the preactivation fields and kernels terms of layer time 7 defined as
T lim E [0, 1]. 1-000 L
(93)
For example, the residual kernels are defined as
Hag (T, x. x', t, t,t')
lim H (x, xt,,t'))
(94)
L-00
The finite difference equations for the residual updates L(h -h²) ~ 37 h(T) become differential updates (either SDE-like ODE-like depending or)) 46 11]
8
E.4.2 ODE Limit oz =
First, we investigate the case of az 1. In this case, the H, L + do limit results in a complete
disappearance of the Xo: Et Ee et fields.
Sq
0,ha(T, 20, 70000 Σ Exrws, A(w' (')V (T, o. w' , 1,1)) l')gg) (7' ,e x',('
Pel
-0. = 1o70 Σ Eg'- A(a', G' (,, t,t') KM (5, t,t')) (2',''e'
33 Σ
Eging A (xx (ax',t'))(Q MM (T, x, x L, Igo x
(95)
C't
45

---

E.4.3 SDE Limit a1
This or, = limit more technically involved since neither the Gaussian terms from DMFT the response functions vanish. For the Gaussian terms, we note that the sums of the independent Gaussians are all multiplied by VL ddr, which can be interpreted integrated Brownian motion in the limit
Σ uh + S du(r')
10000 VL (du(F))du(T)) V - -')
(96)
Following the derivation of Bordelon et al. [II] which maintains the exact dependence on the full integrated response and provides the result as an integrated SDE
he(T, x. t) 30 ! dū,(T' + 7o
Σ PCc
EeiwB, // (T, t't,t))) ( x',t')
(T, x.
R (T, xz x', 1, 1)) + pe (a') (')e
(97)
To To
Combining the forward pass equations from the previous two subsection recovers the Result 3 of the main text. This is combined with complementary equation for the backward pass.
E.5 Effect of MLP Layers
Adding the MIP 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
Bo Laz
hg 1) = h²(x), 1) +
MHSA (x( ()),
(98)
Bo
h (x, h,(x,, t) +
MLP (x.(tt)))
(99)
L
where the MLP layer is
MLP(h VNH
wii o [h.
h
Wiih
(100)
VNH
The following gradient fields necessary
(x, = 7 NH 0f(x,t) t) (2,
0f t) (x,
(x = N7
(x,() No 0f t) 0h 1.1 (x,
(101)
(t2(()) w0(((0) + Bonono I I1-az NNH
EanB Σ A(x, (x,
P't
Bono
(( W + Li-ai VNH
E
Σ
(102)
Pitt
The MLP hidden layer dynamics is much simpler to characterize and resembles the structure analyzed in prior works on infinite width networks [[5].
Bo
h.l (ax, xf. (ax,) 1)
Σ FERB
Ea (a. H (x, x x',t,t) (103)
I1-Oz
P't
Again, we see that the inner dynamics for h. due to the weight updates in this layer scale as T, 7, suggesting the need to choose as we desire this hidden layer to contribute to the representational updates.
46

---

MLP Layer Gradients
For the MLP Mayer we have the simpler backpropagation equations
Oh 0 (x,
(x, ()
9s 5 (x,
oh 4.1 (x,
Bo
d(h (x,, ()) o
wL?(1) 9s g6-1 (xc,
T.
LNNH
(,, = VNH
W (1) g (m,
(104)
The components of these fields that depend initial conditions
1
< 1)
W e. (0) g (x,
VNH
1
(xr, VNH
W (0) (a,
(105)
MLP Matrices After utilizing these resolutions of the identity for all s, x, 4, we can integrate over the weights (c)
Σ do TrW x x t)((h? (x, f)) + go f (x, t) (zx, (7)
(0)) exp
VNH
Σ
doodæ' [xe set (x,t((x,,) (x, 26+1 (x' (x', 1') (ax, x (t,,') + E (x, E (æ', (a=,,")GG (t) x' xt,,)))
1t'
-i Σ dwde' (x6 (x, - (æ', E') E') (x, x' ,t,t))
(106)
tt'ag'
where we introduced the response function
Re (æ, x' ,,,t) NH & (h (x, S. (x',')
(107)
We can perform an identical step integrate over W This gives us
Σ/ da TT) & (a, (,, 1)T (m, ((m.)))( (E((.)))))
In Eoo exp
VNH
Σ
dacdæ' (x,, (x) 261 (xd, (1) () (a, ,',,,,,, ,,,) + ( 1) (dd,("))" (x, x 20',11)) .1,1))
-
1t'ss'
-i E
dwdw' (x, 1) (x', (') (x, x ,1,,)]
(108)
t''a''
where we introduced
Reg (2, x', (,t) NH h,(x,t) t) (²',))
(109)
E.6 Effect of Layer Norm on the Limiting Process
Oh',
Layernorm
The derivative of layer-norm
acts as the following in the large limit ahe
Oh 1
NH 11 1 NH (o2 +c)3/2 h u1l Ih w1]
(110)
0h v0
In the limit of 7/ the variables μ NH h 1 and 0² No h 41 will become deterministic over random initializations. We thus just have to consider how these types of vectors act on gradients
8h
1
0h
9
Vo=+c (g 1n₂)
(o2 + +23022 ,1]
NH
Ih ul] 9)
(111)
Each of these operations will lead to one inner product that will be self averaging M9 N' 1 g or NH hh - w1] g). Thus this operation will not alter the backward pass terms of scaling.
47

---

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 hour to 6 hours depending on model size and total amount of training steps.
The language model trained C4 used context length of 256 and the GPT-tokenizer from Hug- gingface. 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
48

---

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 short (1-2 sentence) justifiiaton 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 (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 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 binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write 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
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We claim to analyze various limits of transformer training and provide theoret- ical and empirical results that effect.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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]
49

---

Justification:
We provide limitations and future directions section in our conclusion.
Guidelines:
The answer NA means that the paper has 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 separate "I "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 sellispecifiitiinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texi 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 cross- referenced.
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 provide 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 properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide example FLAX code which shows how our models are imple- mented. We also mention the datasets and hyperparamcters used.
Guidelines:
50

---

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 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., 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 instruc- tions 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 NeurlPS code and data submission guidelines (https://nips.. cc/
stteeeesubmmsiioooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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 NeurlPS code and data submission guidelines (https
//ipss.cc/publi//guue
Codeeuumiiii wettsbbmmssiiioi ssionoolicy 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 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).
51

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 the core of the paper to 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi- dence 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call 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.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality 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 range (e.g. negative eтror 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 the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide 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 provided in Appendix
Guidelines:
52

---

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 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 failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration 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 of a theoretical nature and is about general scientific understanding of transformer models.
Guidelines:
The answer NA means that there 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 tied to particular applications, let alone deployments. However, if there is direct path to any negative applications,
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, 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 unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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 models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
53

---

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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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 bataasets.ppprsstttt paperswi theode. com/datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of 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 available online, the authors are encouraged reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
54

---

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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, 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 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 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 the institution conducting the review.
55


Reviews:
[{"review_id": "w9keGdKRki", "reviewer": "Reviewer_FJYA", "summary": "The 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.", "questions": "Please see the Weaknesses section.", "limitations": "I 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- 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", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and for allowing us the chance to clarify our theoretical methods. We hope that upon implementing these proposed changes the reviewer will consider improving their score. \n\n### Strengths\nWe appreciate these comments on these strengths of our contributions!\n\n### Weaknesses\n\n### Exposition and Clarity\n\nWe thank the reviewer for this bit of feedback. We have updated the draft to include a larger amount of signposting and exposition (blue text is what has been added since initial submission). \n\n\n### Missing Citations to Non-Asymptotic Approaches\n\nWe thank the reviewer for bringing this body of work to our attention and pointing out its absence in our paper. We now add the following to the related works \n\n\"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 \\cite{bernstein2020distance}. This approach coincides with $\\mu$P width-scaling when the spectral norm of the weights is used as the measure of distance \\cite{yang2023spectral}, and can achieve hyperparameter transfer for a wide array of optimizers and initialization schemes \\cite{bernstein2023automatic, large2024scalable}.\"\n\nWe also point out that Large et al argue for $1/L$ depth scaling, which is similar to our conclusion of the way to scale up depth for transformers. \n\n### Technical Tools and Path Integral Approach\n\n*A major question I have about the work is \"are path integrals necessary here?\"... Why do you use them?*\n\nThis is a great question! \n\nIf one is mainly interested in arguing about whether feature and NN logit updates are $\\Theta_{N,L,\\mathcal H}(1)$ under a given parameterization, then the path integral approach is not necessary. However, if one wants to characterize the exact limits we consider when initializing randomly, then the DMFT method is useful. However, even to derive the limit the path integral method is not the only way to obtain the correct limits (the **dynamical cavity** method also works). \n\nIn response, we have added a new Appendix section D where we explained more clearly what the path integral is computing by giving some simpler examples (see the response section \"Simple Path Integral Examples\"). We will also provide a companion derivation using the **dynamical cavity method** in a simpler setting (generic ResNet with no MHSA blocks) and show that it agrees with the path integral computation. This is very similar to the computation of [Bordelon et al '24](https://arxiv.org/abs/2309.16620) on infinite depth ResNets. Below, we will provide more details about how the path integral is non-vacuous and more systematic than the cavity approach. \n\n### DMFT Path Integral is Non-Vacuous \nDMFT is useful primarily in settings where there is a source of randomness in a high-dimensional dynamical system that *correlates state vectors across time*. The path integral approach gives one the exact asympotic description of the limiting dynamics and can also give finite size corrections. In our problem of interest (training dynamics of transformers) the randomness comes from the random initial weights in each layer, while the states are the features computed on forward and backward passes. In this setting the use of path integrals is far from vacuous (like the $\\ddot x = 0$ example). \n\nThis formalism is very flexible and when the dynamics correspond to some kind of optimization procedure it can be viewed as an alternative to other disordered systems methods (replica method, etc) which *respects the choice of initialization and optimizer*. Some recent examples are studying SGD/momentum with [random data on general linear models](https://arxiv.org/abs/2006.06098) or the [training dynamics of random feature models](https://arxiv.org/abs/2402.01092). We will provide a couple very simple examples that gives intuition for what this approach is computing (see the comment below). It can also be used for problems where there is no equilibrium distribution (like [random RNNs](https://arxiv.org/abs/1809.06042)) or [non-gradient descent learning rules](https://arxiv.org/abs/2210.02157).  \n\n\n#### Cavity Method Alternative\nWhile the path integral approach gets the correct answer, there is an alternative **cavity method** derivation that provides a different set of intuitions while recovering the same limiting dynamics. We will include a simple cavity derivation in Appendix D. The idea of the cavity method is to consider what happens when a single neuron is added to the residual stream or one of the hidden layers. This new neuron gives small corrections to all other neurons, but these add up to give the extra response terms. \n\n\n#### Merits of Path Integral Compared to Cavity Method\n1. It starts by giving the non-asymptotic distribution of the feature kernels and response functions at finite $N$ induced by randomly sampling initial weights. This distribution looks like $p(Q) = \\frac{1}{Z} \\exp\\left( - N S[Q]  \\right)$ where $S$ is the DMFT action and $Q = \\{f , \\Phi, G, R, ... \\}$ are the network outputs $f$, feature kernels $\\Phi$, gradient kernels $G$ and response functions $R$ (everything that should concentrate as $N \\to \\infty$).\n2. While the $N \\to \\infty$ limit is computed as a saddle point (first derivative) of the DMFT action $\\frac{\\partial S}{\\partial Q} = 0$), finite width corrections can also be obtained from [higher order derivatives](). One can track the leading order $\\mathcal{O}(N^{-1})$ corrections to the dynamics of $Q$ from the Hessian $\\frac{\\partial^2 S}{\\partial Q \\partial Q}$. \n3. The cavity method often requires \"having a sense of the final result\" before doing the computation. The computation is easy if you already possess intuition for the kind of mean-field limit you expect to get and what response functions will appear. The path integral method is more systematic and requires less mean field theory \"intuition\".", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OydGaEtUk2", "reviewer": "Reviewer_rXAV", "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:\n\n- Under the limit $N\\to\\infty$ the $\\mu P$ rule is required, causing all heads to collapse into the same dynamics.\n- Under the limit $H\\to\\infty$, each head becomes statistically independent.\n- The scaling under the depth limit $L\\to\\infty$ is also analyzed.", "questions": "- Line 78: What is $\\gamma_0$? Could you elaborate on this further?\n- What is the maximum model size (i.e., number of parameter) used in the experiments?\n- 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.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n- 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).\n- The following paper is also relevant:\n\n  - Lénaïc Chizat, Praneeth Netrapalli. The Feature Speed Formula: A flexible approach to scale hyper-parameters of deep neural networks. 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Strengths\n\nWe thank the reviewer for appreciating these aspects of our contributions.  \n\n### Weaknesses\n\nBelow we try to clarify our limits and also add a citation to the \"Feature-Speed Formula\" paper and discuss the similarities and differences of the conclusions we make in this paper. We hope that in response the reviewer would be willing to increase their score. \n\n### Practical Utility of Studying Infinite Head Transformers\n\nThis is a good question. To motivate the infinite head limits we point out the following few facts. \n\n1. We first investigated the large key-query dimension $N \\to \\infty$ limit of transformers with $\\mathcal H$ fixed. We showed that this limit was degenerate in the sense that all attention heads collapsed to the same dynamics, causing redundant computation in a multi-head model. We thus sought another way to take \"width\" to infinity that would allow variability across attention heads.    \n2. In the scaling law era, model sizes keep increasing and often heads $\\mathcal H$ are increased as well as depth $L$ and key/query dimension $N$ as the models are scaled up. In the [GPT-3 technical report](https://arxiv.org/abs/2005.14165) Table 2.1, we can see that the last few models are increased by scaling up heads and layer count with $N$ (they call $d_{head}$) fixed. Understanding if this scaling behavior leads to a well defined limit is therefore an interesting and potentially practical question to guarantee stable convergence.  \n3. Even if models are finite, they often share important similarities with their larger (or infinite) width counterparts (for mean field/$\\mu$P scalings) and thinking about limiting behavior can be useful when designing parameterizations for width/depth. For example, optimal learning rates often transfer across models of different [widths](https://arxiv.org/abs/2203.03466) and [depths](https://arxiv.org/abs/2309.16620) when a scaling limit exists. Models [learn similar representations](https://arxiv.org/abs/2305.18411) across different model sizes. Depending on the setting or number of training steps, the infinite limit can be very descriptive of models with widths as modest as a few hundred. For Figure 4(a) we show that $\\mathcal H = 16$ and $\\mathcal H = 128$ models can have very similar training dynamics on CIFAR-5M.  \n4. Developing mean field theory for infinite limits (specifically the DMFT action) can enable one to obtain the [dynamics of finite size corrections](https://arxiv.org/abs/2304.03408) to the theory. These would capture an approximate evolution of the $\\frac{1}{\\sqrt{\\mathcal H}}$ deviation from the infinite head limit if one finds the infinite head limit unrealistic.\n\n### Feature-Speed Formula\n\nThank you for pointing out this interesting paper.  The authors of this work develop theory for deep network Jacobians at initialization and early feature learning updates. We have added a citation to this work in the related works section. \n\n#### Desiderata of Chizat and Netrapalli\nThey point out four desiderata of a scaling limit to be the following four criteria at initialization\n1. Signal propagation\n2. Feature learning \n3. Loss Decay\n4. Balanced Contributions across layers\n\nThe authors analyze conditions under which these four conditions hold in large depth networks at initialization. They conclude that residual networks $(h^\\ell+1 = h^\\ell + \\beta W^\\ell \\phi(h^\\ell)$) with branch scale $\\beta = 1/\\sqrt L$ are necessary.\n\n#### Why are Our Conclusions about Depth Different?\nSome differences between our work and this work's conclusions\n1. We consider transformer models where there are multiple layers per residual block. In these models, if the residual scale factor is $\\beta = 1/\\sqrt{L}$ then the hidden key/query weights $W_K^\\ell, W_Q^\\ell$ in the attention layer can be treated as frozen in the $L \\to \\infty$ limit. However, under the $\\beta = 1/L$ branch scaling, these matrices update non-negligbly, leading to additional feature learning (see Figure 1(c) in Rebuttal pdf). \n2. While each layer's initial contribution to the initial neural tangent kernel is not balanced if $\\beta = 1/L$, over the course of training they will become balanced. \n\nIn the related works section, we add a citation to Chizat and Netripalli and also include the sentence\n\n\"In this work, we pursue large depth limits of transformers by scaling the residual branch as $L^{-\\alpha_{L}}$ with $\\alpha_{L} \\in [\\frac{1}{2},1]$ ... However, we argue that in transformers that $\\alpha_L = 1$ is preferable as it enables the attention layers to update non-negligibly as $L \\to \\infty$.\"\n\n\n### Response to Questions\n\n*Line 78: What is $\\gamma_0$? Could you elaborate on this further?*\n\nThe parameter $\\gamma_0$ is a constant $\\Theta_{N,H,L}(1)$ hyperparameter that controls the scale of hidden feature updates ([laziness/richness](https://papers.nips.cc/paper_files/paper/2019/hash/ae614c557843b1df326cb29c57225459-Abstract.html)). The $\\gamma_0 \\to 0$ limit gives a lazy learning limit. We borrow this notation from previous papers on the DMFT limits.\n\nTo make this more clear we give the equation that defines $f$ its own line \n\n\\begin{align}\n    f = \\frac{1}{ \\gamma_0 N \\mathcal H} w^L  \\cdot \\left( \\frac{1}{\\mathcal S} \\sum_{\\mathfrak s}  h^L_{\\mathfrak s} \\right)\n\\end{align}\n\n*What is the maximum model size (i.e., number of parameter) used in the experiments?*\n\nThe maximum number of parameters in our CIFAR-5M plots are around $100$M parameters (which was the $\\mathcal H = 2048$ model in Figure 3) while for the C4 language experiments the maximum model size was $150$M parameters. \n\n*Do the results hold for any time t? Prior works focus on the initialization phase.*\n\nWe are also explicitly keeping timesteps fixed as we scale the model size (not scaling these jointly). This was mentioned in a footnote and in the limitations section, but we added an extra sentence in the main text to emphasize this.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RKduvlCZd4", "reviewer": "Reviewer_CMzU", "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.", "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)\n2. 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?\n3. Why 3/2 appears in eq3\n4. 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", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Give detailed analysis (and closed form) on dynamics of the updates\n- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Strengths\n\nWe thank the reviewer for appreciating these aspects of our paper and for their support.\n\n### Weaknesses\n*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.*\n\nThank you for this advice. We will now provide more detail in the main text about how the DMFT works and also add a new Appendix section which gives a primer on the DMFT methods used in this work.\n\n### Questions:\n\n*In line 78, it is mentioned that gamma_0 controls the rate of feature learning, but I couldn't find it in equation (1)*\n\nThe parameter $\\gamma_0$ is a constant $\\Theta_{N,H,L}(1)$ hyperparameter that controls the scale of hidden feature updates in the same way described by Chizat and Bach. The $\\gamma_0 \\to 0$ limit gives a lazy learning limit. We borrow this notation from previous papers on the DMFT limits.\n\nTo make this more clear we give the equation that defines $f$ its own line \n\n\\begin{align}\n    f = \\frac{1}{ \\gamma_0 N \\mathcal H} w^L  \\cdot \\left( \\frac{1}{\\mathcal S} \\sum_{\\mathfrak s}  h^L_{\\mathfrak s} \\right)\n\\end{align}\n\n\n*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?*\n\nYes, this phenomenon is due to the fact that Adam updates are approximately normalized. Consider a simple MLP layer. We would have a change to the weight of the size $\\delta W_{ij} \\sim \\eta \\frac{g_{i} \\phi_j}{\\sqrt{g_i^2 \\phi_j^2 + \\epsilon}}$. We need to control the size of the forward pass\n\\begin{align}\n    \\frac{1}{\\sqrt N } \\delta W \\phi(h) \\sim \\Theta( \\eta N^{1/2} )\n\\end{align}\nWe want this to be $\\Theta(1)$ so we need $\\eta \\sim N^{-1/2}$. If you desire more details, check out Appendix C.3.\n\n*Why 3/2 appears in eq3?*\n\nThe $3/2$ is due to the fact that feature updates have to be controlled to prevent $k \\cdot q / N^{\\alpha}$ from diverging. We work this out in Appendix C.2.\n\n*Why the limit can be straightforwardly computed from the saddle point of DMFT action?*\n\nWe have added a more detailed Appendix section on DMFT which shows how the path integral method works in more detail on a few simple problems. In a nutshell though, the DMFT approach shows that the distribution of the feature kernels and response functions at finite $N$ induced by randomly sampling initial weights looks like $p(Q) \\sim \\exp\\left( - N S[Q]  \\right)$ where $S$ is the DMFT action and $Q = \\{f , \\Phi, G, R, ... \\}$ are the network outputs $f$, feature kernels $\\Phi$, gradient kernels $G$ and response functions $R$ (everything that should concentrate as $N \\to \\infty$). By a steepest descent argument, the probability density $p(Q)$ will be dominated at large $N$ by the saddle point. This is a useful idea in mean field theory and high dimensional statistics.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z2z1zngaoS", "reviewer": "Reviewer_UGyb", "summary": "The 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.", "questions": "* 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?", "limitations": "A Limitations section is provided in Section 5.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "While 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.", "weaknesses": "See Questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Strengths\n\nWe thank the reviewer for their supportive comments and for their detailed reading of our work. Below we try addressing the questions and mentioning ways we aim to improve the paper.  \n\n### Questions\n*1. 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\nThis is a good question! The central limit theorem and law of large numbers do give you pretty good intuition in general and motivate the choice of scaling exponents. However, to make fine grained conclusions about the exact limiting dynamics one often needs to be careful with extra terms that appear in the dynamics known as response functions which are not apriori obvious (see comment below on non-trivial DMFT examples). \n\nOne precise insight from our analysis is the symmetry across heads of the resulting limiting equations as key query dimension diverges $N \\to \\infty$. This led to our conclusion that heads degenerate into the same dynamics and motivated us to characterize the infinite head $\\mathcal H \\to \\infty$ limit. We can also give some similar qualitative insights about the types of dynamics one obtains as $L \\to \\infty$ from our equations (for instance response functions are negligible for $\\alpha_L = 1$ but survive for $\\alpha_L = \\frac{1}{2}$). However, the exact equations in the general case are complicated non-Gaussian stochastic processes. In **linear** transformers, we suspect these tools can give much more interpretable insights, but this case is not as realistic. \n\n*2. 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\nWe have added some preliminary comparisions of this kind. It does seem that in terms of width scaling, increasing $\\mathcal H$ is preferable to increasing $N$ with $\\alpha_{\\mathcal A} = 1$. In addtion, when scaling depth $L$ it appears that $\\alpha_L = 1$ is preferable to $\\alpha_L = \\frac{1}{2}$, consistent with our theory. \n\n*3. The distinction between $\\alpha$ and is explained in Appendix C.4, what happens if interpolates between the two extremes?*\n\nFor $\\alpha_{\\mathcal A}$ between $\\frac{1}{2}$ and $1$, the the variables $\\mathcal A = k \\cdot q / N^{\\alpha}$ will concentrate to zero at initialization for any $\\alpha > \\frac{1}{2}$ leading to the head-collapse issue as $N \\to \\infty$. If one tunes learning rates to make $\\mathcal A$ change substantially with SGD, it will lead to a divergence in $N$ unless $\\alpha < 1$. Thus $\\alpha = 1$ is special as it enables stability as $N \\to \\infty$ and $\\alpha = \\frac{1}{2}$ is special as it preserves $\\Theta(1)$ diversity of heads at initialization. \n\n*4. What are some previous results or potential approaches towards analyzing limiting dynamics of neural networks under Adam?*\n\nSome recent works have attempted to investigate the limiting behavior of Adam using Tensor programs https://arxiv.org/abs/2308.01814 or from an empirical scaling perspective https://openreview.net/pdf/579c102a8c067102c85e27612c36d7a356ea9b0b.pdf. One insight into these analyses is that the $\\epsilon$ parameter in Adam may have to be scaled with width to obtain a reasonable limit.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Global Response

We thank the reviewers for all of their detailed comments and advice on ways to improve the paper. Below we go through some of the concerns which arose in comments from many reviewers and outline how we plan to address them in the newer version of the draft. 

### Repeated Concerns

1. Path integrals / DMFT / Saddle points may not be familiar to the ML audience so some additional exposition about these would be useful. We also point out that deriving the DMFT action (at least in principle) gives a procedure to also extract finite size effects.
2. What is the motivation for taking infinite limits?
3. Exposition and labeling: we have tried making the paper more readable with additional signposting and better labeling on Figures and Figure captions. Some variables (such as $\gamma_0$) are not well motivated or defined. 
4. Some missing citations to relevant prior works.
5. Request for additional clarification of the limitations of our theory. 
6. Compute optimal comparisions: how do these parameterizations or scaling limits compare as compute is varied?

### Updates to the Paper in Response

In response to these issues, we will make the following updates which will appear in any future version of the paper.

1. We have added a short expository section in the main text and a new Appendix section which gives a primer on DMFT methods and what the path integral approach is computing. In addition, we provide simple but qualitatively similar examples to deep network dynamics where the path integral gives the correct limiting stochastic dynamics of hidden neurons. We also now provide a companion derivation of the limit using the **dynamical cavity method** showing a physical interpretation of the response functions.
2. We try motivating infinite limits since (1) models improve in performance as parameter count increases (2) infinite limits can be descriptive of finite models (3) theory for infinite models can often be extended to approximate finite models.
3. We aim to clean up the exposition. Before using a letter $N,\mathcal{H},L$ we preempt with "number of attention heads $\mathcal{H}$". We also will fix some legibility and colorbar issues with our plots.  We now explicitly define the feature learning scale $\gamma_0$ which is controls the [laziness/richness](https://papers.nips.cc/paper_files/paper/2019/hash/ae614c557843b1df326cb29c57225459-Abstract.html) of the training. This notation is adopted following [prior works on mean field limits](https://iopscience.iop.org/article/10.1088/1742-5468/ad01b0/meta). 
4. We have added citations to works on non-asymptotic approaches to stable training across widths and depths (including Bernstein et al 2021 and the Modula paper from Large et al 2024) which have an interesting alternative non-asymptotic perspective on hyperparameter transfer and width/depth scaling. We also now cite and discuss the "Feature-speed formula" paper from Chizat and Netrapalli which discusses the stable $\frac{1}{\sqrt L}$ branch scaling in residual networks in the feature learning regime. See the Attached Figure 1c in rebuttal for information about why $\alpha_{L}=1$ uniquely allows attention layers to update. 
5. We now emphasize that the derived theory holds for fixed training horizons (training time is treated as a constant that is not scaling jointly with width/heads/depth). As reviewer FJYA points out, finite size effects from the stochastic initialization can accumulate over time.
6. We provide some plots of performance as a function of compute in our experiments in the attached rebuttal document. We find that scaling $\mathcal H$ is preferable to scaling $N$ in the parameterizations that admit suitable limits. In addition, we also find that $\alpha_L = 1$ is preferable to $\alpha_L = \frac{1}{2}$, consistent with our theory. We plan to add these and more experiments like these to the paper. 

### Added Expository paragraph on DMFT in Main Text

"To obtain the exact infinite limits of interest when scaling dimension-per-head $N$, the number of heads $\mathcal H$, or the depth $L$ to infinity, we work with a tool from 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 data. We use this method to reason about the dynamics of randomly initialized neural networks by tracking a set of deterministic correlation functions (feature and gradient kernels) as well as response functions (see Appendix D). The core conceptual idea of this method is that in the infinite limit, all neurons remain statistically independent throughout training and only interact through collective variables (feature kernels, neural network outputs, etc). Collective variables are *averages* over the distribution of neurons in each hidden layer or along the residual stream."

### Added Sentence about Fixed Training Time 

At the bottom of page 2 we added the sentence
"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."  

### New Appendix Section 

Our new Appendix D "Primer on DMFT" will contain more detailed information about DMFT and the path integral method. We motivate DMFT as a general technique for dealing with dynamical systems that depend on **fixed sources of randomness**. We provide a few simple examples where one can see that the resulting stochastic process is non-trivial including

1. A linear dynamical system driven by a random matrix
2. Feature updates in deep linear neural networks and residual networks.

In both of these settings, the DMFT linear-response functions give non-trivial corrections to the limiting dynamics. 

We are also adding information about the alternative **dynamical cavity** method to derive the DMFT equations which does not require the use of path integrals.

Meta-review: This paper provides various scaling limit of multi-head attention in the feature learning regime as $N,H,L \to \infty$. In the limit of $N \to \infty$, the $\mu P$ parameterization is required, and the DMFT is applied to derive the training dynamics in the limit. Some numerical experiments are provided to show their convergence.  

The analysis provides a precise description of the infinite limit of $N,H,L$. Although we need a careful setting of parameter scaling to achieve feature learning, this paper provides a comprehensive theoretical investigation to this issue. 
Overall, the paper's technique is novel. Its contribution would be appreciated by the community.   

As a summary, I think this paper is worth publication in NeurIPS. On the other hand, I recommend the authors to proofread the paper again.
Decision: Accept (poster)