Infinite Limits of Multi-head Transformer Dynamics

Authors: blake  bordelon
hamza  chaudhry
cengiz  pehlevan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 00:42:39

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.

Main paper: 
 introduction: 
 Increasing the scale of transformer models has continued to improve performance of deep learning systems across many settings including computer vision [1, 2, 3, 4] and language modeling [5, 6, 7, 8, 9] . However, understanding the optimization stability and limiting behavior of these models under increases in model scale remains a core challenge.One approach to scaling up systems in a stable and predictable way is to identify parameterizations of neural networks that give approximately scale-independent feature updates during training [10, 11, 12] . The mean field parameterization, commonly referred to as µP, is a well-known example that satisfies this property [13, 14, 15] . When such parameterizations are adopted, the learned internal representations in hidden layers of the network are very similar across model scales [16, 17] , but performance tends to improve with model scale [10, 11, 12] . Further, theoretical results about their limits can often be obtained using Tensor Programs [14] or dynamical mean field theory (DMFT) techniques [15, 17] .In this work, we develop a theoretical treatment of randomly initialized transformers. We study various scaling limits of the training dynamics of these models including the infinite key/query dimension limit, the infinite head limit, and the infinite depth limit. Concretely, our contributions are the following:1. We derive a DMFT for feature learning in randomly initialized transformers with key/query dimension N , attention head count H and depth L. From the derived DMFT action, we identify large N , large H and large L limits of the training dynamics.2. We analytically show that the large key-query N → ∞ limit requires the µP scaling of key/query inner product with 1/N , even if key/queries are reparameterized to decrease the size of their updates from gradient descent.3. From the limiting equations, we show that this N → ∞ limit causes multi-head self attention trained with stochastic gradient descent (SGD) to effectively collapse to single-head self attention since all heads follow identical dynamics.4. To overcome this limitation, we analyze the infinite head H → ∞ limit while fixing N . We show there is a limiting distribution of attention variables across heads at each layer throughout training. Despite N being finite, the infinite-head H → ∞ limit leads to concentration of the network's output logits and learned residual stream feature kernels, giving deterministic training dynamics.5. Finally, we examine large depth limits of transformers with residual branch scaling. We illustrate and discuss the tension between parameterizing a model so that it has a non-trivial kernel at initialization while maintaining feature learning within the multi-head self attention (MHSA) and multi-layer perceptron (MLP) blocks.Hron et al. [18] studied the Neural Network Gaussian Process limit of multi-head self attention in the infinite-head H → ∞ limit. They showed that, at initialization, there is a limiting distribution over attention matrices and that the outputs of the multi-head attention block follow a Gaussian process, establishing a connection to kernel methods. Dinan et al. [19] develop a similar theory of transformers at initialization and compute the Neural Tangent Kernel associated with this architecture as the dimensions per head N → ∞ using a 1 √ N scaling of the key-query inner product within each attention layer. One of our key theoretical results is showing that this picture of a distribution over learned attention heads persists throughout training in the feature-learning regime as H → ∞ (though the distribution of residual stream variables generally becomes non-Gaussian).Several works have analyzed the signal propagation properties of transformers at initialization at large key/query dimension N and large depth L [20, 21, 22, 23] including providing modifications to the standard transformer architecture [22, 24] . In this work, we pursue large depth limits of transformers by scaling the residual branch as L -α L with α L ∈ [ 1 2 , 1], which has been shown to converge to a limit not only at initialization [25, 26, 27] , but also throughout training in the feature learning regime [11, 12, 27] . However, we argue that in transformers that α L = 1 is preferable as it enables the attention layers to update non-negligibly as L → ∞.Yang et al. [10] introduced the µP scaling for attention layers which multiplies the key/query inner product with 1 N rather than the more commonly used 1 √ N [5] . They show empirically that this change improves stability of training and transfer of optimal hyperparameters across different values of N . Vyas et al. [16] empirically found that such µP transformers learn attention matrices that become approximately consistent across different heads and model sizes, suggesting that models parameterized in µP learn similar representations across scales.In addition to work on infinite width and depth limits of deep networks, there is also a non-asymptotic approach to optimizer design and scaling based on controlling the norm of weight updates [28] . This approach coincides with µP width-scaling when the spectral norm of the weights is used as the measure of distance [29] , and can achieve hyperparameter transfer for a wide array of optimizers and initialization schemes [30? ]. 
parameterizations with feature learning limits: 
 We consider a transformer architecture with L layers, H heads per layer, and N dimensional keys/ queries per head. Transformers are often defined in terms of d model = Hd head = HN which can be increased by scaling the number of heads or the dimension of each head, where N is often written d head . Our goal is to determine the set of parameterizations that allow the attention layers to undergo non-trivial feature learning in the various N, H, L → ∞ limits. The analysis of these limits is performed with batch size and number of training steps t fixed while the other architectural parameters are taken to infinity.The network's output is computed by a depth L recursion through hidden layers ℓ ∈ [L] starting with the first layer h 1 s (x) = 1 √ D W 0 x s ∈ R N H where x s ∈ R D is the input at spatial/token position s. Preactivations in subsequent layers h ℓ are determined by a forward pass through the residual stream which contains an attention layer and a MLP layerEQUATIONThe constants γ 0 and β 0 control the rate of feature learning and the effective depth respectively 1 . We will consider α L ∈ [ 1 2 , 1]. 2 The multi-head self attention layer (MHSA) with pre-layer-norm 3 isEQUATIONwhere σ ℓ h ∈ R S×S are the attention matrices passed through a matrix-valued nonlinearity σ A ℓ h 4 .For a sequence of length S, the pre-attention matrix A ℓ h ∈ R S×S is defined asEQUATIONThe exponent α A will alter the scale of the pre-attention variables A ℓ h at initialization. The input matrices have shape W ℓ V h , W ℓ Kh , W ℓ Qh ∈ R N ×N H , while the output matrices have shape W ℓ Oh ∈ R N H×N . All of the weights W ℓ Oh ′ , W ℓ Qh , W ℓ Kh are initialized with Θ(1) entries while W ℓ Kh , W ℓ Qh have entries of size Θ(N 1-α A ) which ensures that all key and query k, q vectors are Θ(1) at initialization. The pre-attention variables A ℓ h ∈ R S×S at each head h are determined by key k ℓ hs and query q ℓ hs ′ inner products. The MLP layer consists of two linear layers with an element-wise nonlinearity ϕ applied in between, where W ℓ,2 , W ℓ,1 ∈ R N H×N H are initialized with Θ(1) entries:EQUATIONµP scaling [13, 31, 14, 15] downscales the readout of the last layer compared to standard and NTK parameterization [32] . Thus, we define the output of the model asEQUATION5 where h L s ∈ R N H are the final layer preactivations at spatial/token position s ∈ [S]. The parameter γ 0 is an additional scalar that controls the rate of change of the internal features of the network relative to the network output [33] .In order to approximately preserve the size of internal feature updates, we must scale the learning rate η appropriately with (N, H, L). However, this scaling depends on the optimizer. In Table 2 , we provide the appropriate scaling of learning rates for SGD and Adam for any α L ∈ [ 1 2 , 1] andα A ∈ [ 1 2 , 1].In what follows, we focus on the SGD scaling and theoretically analyze the N → ∞, H → ∞, and L → ∞ limits of the training dynamics. We also provide in Table 2 details about what prefactor the first layer should be multiplied by and the initial weights divided by to ensure convergence to the L → ∞ limit. Example FLAX implementations of these parameterizations for vision and language modeling transformers are provided in Appendix B. Global LR First/Last Layer Rescale Multiplier First/Last Layer Std. Dev. SGDη 0 N HL 2α L -1 L 1 2 -α L Θ(1) Adam η0 √ N H L -1+α L L 1-α L √ N H 1 √ N H L -1+α L Table 1:The learning rates which should be 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 (relative to eq (5) ) with width/depth depending on the parameterization and optimizer.Our analysis assumes that at each step t of SGD or Adam a mini-batch B t of size Θ(1) is used to estimate the loss gradient. We assume that the minibatches are fixed. Further, the number of total training steps is assumed to not be scaled jointly with the model size. Therefore the analysis provided here can cover both online SGD for a fixed number of steps or full batch GD (repeating data) with a Θ(1) sized dataset. 
infinite limits of learning dynamics: 
 In this section, we first analyze the infinite dimension-per-head N → ∞ limit of training. We find that for this limit, the µP rule of α A = 1 is necessary and show that all heads collapse to the same dynamics. To counteract this effect, we next analyze the infinite head H → ∞ limit of the training dynamics at fixed N, L, where we find a limiting distribution over attention heads. We will conclude by analyzing the statistical descriptions of various infinite depth L → ∞ limits. 6 .To obtain the exact infinite limits of interest when scaling dimension-per-head N , the number of heads H, or the depth L to infinty, we work with a tool from statistical physics known as dynamical mean field theory (DMFT). Classically, this method has been used to analyze high dimensional disordered systems such as spin glasses, random recurrent neural networks, or learning algorithms with high dimensional random data [34, 35, 36, 37, 38, 39] . Following [15, 11] , we use this method to reason about the limiting dynamics of randomly initialized neural networks by tracking a set of deterministic correlation functions (feature and gradient kernels) as well as additional linear-response functions (see Appendix D). The core conceptual idea of this method is that in the infinite limit and throughout training, all neurons remain statistically independent and only interact through collective variables (feature kernels, neural network outputs, etc). Further the collective variables can be computed as averages over distribution of neurons in each hidden layer or along the residual stream. This DMFT description can be computed using a path integral method that tracks the moment generating function of the preactivations or with a dynamical cavity method (see Appendix D).One way of obtaining a well-defined infinite parameter limit of transformers is to take the N → ∞ limit, where N is the dimension of each head. A priori, it is unclear if there are multiple ways of scaling the key/query inner product. Concretely, it is unknown what values for the exponent α A are admissible for the pre-attention A = 1 N α A k • q. The keys and queries are uncorrelated at initialization which motivated the original choice of α A = 1 2 [5, 18] . Yang et al. [10] assume the entries of the key and query vectors move by Θ(1), implying α A = 1 is necessary since the update to k is correlated to q and vice versa. However, it is possible to obtain Θ(1) updates to the attention variable for alternative values of α A if we choose the change to key (also query) entries after gradient descent to be δk i ∼ Θ(N -1+α A ). We show that this scaling can approximately preserve optimal hyperparameters across N in Figure 2 (a) and give similar dynamics under SGD Appendix C. However, as we show in Appendix E.1.2, any well defined N → ∞ limit of SGD requires α A = 1.The reason is not that keys and queries become correlated, but rather that the scale of the backward pass must be controlled to ensure the dynamics remain stable (non-divergent) under SGD training.After performing two or more gradient descent steps, we demonstrate that the backpropagation signals will diverge as N → ∞ unless initial key and query weight matrices are downscaled to have variance of order Θ N (1). In Appendix E, we provide a DMFT analysis of the N → ∞ limit of the transformer training dynamics. We summarize the result of that analysis informally below.Result 1 (Infinite Dimension-Per-Head N ) (Informal) A stable feature learning N → ∞ limit of transformer SGD training requires taking α A = 1 (µP scaling), even if key/query updates are allowed to be rescaled to account for their correlation. The limiting dynamics of training are governed by the residual stream kernelH ℓ ss ′ (x, x ′ , t, t ′ ) = 1 N H h ℓ s (x, t) • h ℓ s ′ (x ′ , t ′) and a collection of inner product kernels in each head h that concentrate as N → ∞EQUATIONalongside residual-stream gradient kernels and response functions in the sense of [15, 11] . The NN output logits f (x, t) evolve deterministically according to the above kernels as well as kernels for the gradient vectors g ℓ ≡ γ 0 N H ∂f ∂h ℓ which appear in the backward pass. These variables become identical across heads such that for any h, h′ ∈ [H], A ℓ hss ′ (x, t) = A ℓ h ′ ss ′ (x, t).All preactivations on the residual stream and key/query/value variables within a MHSA block are statistically independent across neurons and can be described by a single scalar stochastic processEQUATIONwhere ũℓ , u ℓ , u ℓ Kh are Gaussian processes with covariances Φ ℓ,1 , V ℓσ , H ℓ respectively. Analogous equations hold for the queries and values. In the limit, the kernelsH ℓ ss ′ (x, x ′ , t, t ′ ) = h ℓ s (x, t)h ℓ s ′ (x ′ , t ′ ) , A ℓ hss ′ (x, t) = k ℓ hs (x, t)q ℓ hs ′ (x, t), etc. are computed as averages ⟨•⟩ over these random variables. The deterministic kernels C ℓ , Cℓ can also be expressed in terms of single site averages of residual variables and head averages of MHSA variables. The kernels C ℓ , Cℓ , C k ℓ depend on the precise mini-batches of data B t presented at each step t which we assume are known.We derive this result using a Martin-Siggia-Rose path integral formalism [40] for DMFT in Appendix E. Full DMFT equations can be found in Appendix E.2. Following prior works on DMFT for infinite width feature learning, the large-N limit can be straightforwardly obtained from a saddle point of the DMFT action [15, 11, 41, 17] .As N → ∞, multi-head self-attention will effectively compute the same outputs as a single-head self-attention block. We theoretically derive this effect in Appendix E.2.1 and demonstrate it empirically in Figure 2 (b). This experiment shows that in α A = 1 (µP) transformers trained for 2500 steps on CIFAR-5M [42] , the variance of attention matrices across heads decreases with N . However, we note that if the scaling exponent is chosen instead asα A = 1 2there is non-decreasing diversity of attention variables across heads. This result is consistent with recent empirical findings that attention variables in µP transformers converge to the same limiting quantities at large N with H fixed for different initializations and also across model sizes [16] . This aspect of transformers in the large-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 if there exist well defined limits at finite N with a diversity of attention variables across heads.In this section, we take H → ∞ with the inner dimension N fixed. Rather than having all kernels concentrate, the kernel of each head of the MHSA block follows a statistically independent stochastic process. This picture was shown to hold at initialization by Hron et al. [18] . Here, using a DMFT analysis, we show that it continues to hold throughout training in the feature learning regime.Result 2 (Infinite Head Limit) (Informal) The H → ∞ limit of SGD training dynamics in a randomly initialized transformer at any key/query dimension N , scaling exponents α A , α L , and any depth L is governed by head-averaged kernels for pairs of input data x, x ′ at training times t, t ′ and spatial/token positions s, s ′ such asEQUATIONwhich converge to deterministic values as H → ∞.The attention variables {k ℓ h (x, t), q ℓ h (x, t), v ℓ h (x, t), A ℓ h (x, t)} within each head become statistically independent across heads and decouple in their dynamics (but not across dimensions within a head). Each residual stream neuron becomes independent and obeys a single site stochastic process analogous to Result 1, but with different kernels.We derive this and the full DMFT in Appendix E.3, showing that the joint distribution of headaveraged dynamical quantities satisfies a large deviation principle and the limit can be derived as a saddle point of a DMFT action.To gain intuition for this result, we first examine variables H ℓ and A ℓ h at initialization. In Figure 3 , we plot the convergence of a N = 4, L = 8 vision transformer's residual stream kernel H ℓ to its H → ∞ limit at rate O(H -1 ) in square error, consistent with perturbative analysis near the limit [17] . Next, we plot the distribution (over heads) of A h at a fixed pair of spatial/token positions for a fixed sample. This is a non-Gaussian random variable for finite N , but as N → ∞ the distribution of A will approach a Gaussian with variance Θ(N 1-2α A ).We then investigate training dynamics as we approach the H → ∞ limit. In Figure 4 (a) we show the test loss on CIFAR-5M as a function of the number of training iterations. The performance tends improve as H increases and the model approaches its limit. In Figure 4 (b) we show that all of the models are converging in function space by measuring the squared error between finite H head models and a proxy for the infinite H model. Since the H → ∞ limit is essentially uncomputable, we approximate it as the ensemble averaged predictor of the widest possible models, a technique used in prior works [16, 11] . We again see that at early time, the logits of H head models converge to the limit proxy at a rate O(H -1 ), but after continued training the convergence rate weakens. This effect 2 ) and losses averaged over 10 random inits (colored error bars are standard deviations). (a) As H increases the loss and the variability over random initial seed decreases. (b) The mean square difference between output logits for H head models and a proxy for the infinite head model on a held out batch of test examples. Following prior works, our proxy for the limit is the ensemble averaged outputs of the widest models [16, 11] . has been observed in µP networks in many settings [16] and a theoretical model of this was provided in recent work which argues it arises from low-rank effects in the finite H kernels [39] .H ℓ ss ′ (x, x ′ ) = 1 HN h ℓ s (x) • h ℓ s ′ (x ′ ) in a L = 8, N =We next describe the infinite depth limits which depend on the choice of α L . Below we informally describe the main finding which again uses a DMFT formalism and is based on analyses in recent works on infinite depth networks from Bordelon et al. [11] and Yang et al. [12] . Result 3 (Infinite Depth Limit) (Informal) The training dynamics for H, L → ∞ with L -α L branch scaling with α L ∈ [ 1 2 , 1] is described by a differential equation for residual variables h s (τ, t) in layer time τ = lim L→∞ ℓ L for the residual streamEQUATION)where the Brownian motion term du s (τ, x, t) survives in the limit only if α L = 1 2 and has covarianceEQUATIONand the deterministic kernel C ss ′ (τ, x, x ′ , t, t ′ ) can be expressed in terms of head-averaged kernels and response functions. The weights inside each hidden MHSA layer or each MLP layer are frozen in the L → ∞ limit unless α L = 1. All response functions are suppressed at L → ∞ unless α L = 1 2 . Below we provide a couple of short comments about this result. The proof and full DMFT is provided in Appendix E.4.1. At initialization t = 0, the only term which contributes to the residual stream layer dynamics is the integrated Brownian motion τ 0 du(τ ′ ) which survives at infinite depth for α L = 1 2 . For α L = 1 this term disappears in the limit. The structure of C(τ ) is also modified by additional response functions at α L = 1 2 [11] which we show disappear for α L = 1. 2. The weights within residual blocks (including the MHSA block) can be treated as completely frozen for α L < 1 in the L → ∞ limit, which leads to the simplified statistical description of the preactivations in those layers. However, the residual stream variables h(τ ) do still obtain Θ L (1) updates. At α L = 1 the weights in the MHSA blocks evolve by Θ L (1), causing additional feature evolution in the model.3. A consequence of our large N and large L result is that the N, L → ∞ limit with α L = 1 2 (the parameterization studied by [11, 12] ) would lead to A ℓ hss ′ (x, t) = 0 for all time t. Thus the MHSA blocks would only involve average pooling operations over the spatial indices, despite the residual stream kernels H ℓ updating from feature learning. Test Loss The compute scaling laws with models at fixed width N, H and varying depth L. At large L, the α L = 1 (dashed) models perform better at fixed compute.L = 1 2 L = 1 L = 4 L = 8 L = 16 L = 32 (b) Depth Scaling for αL ∈ { 1 2 , 1}First, we note in Figure 5 that the weights within each attention block freeze as L → ∞ with α L = 1 2 case but move at a constant scale for α L = 1. As a consequence, the loss at large L can be lower in the α = 1 parameterization. We can see some numerical evidence for the first of these effects in Figure 6 (a)-(b) where initially training at large L is slower than the base model and the initial kernel appears quite different for L = 4 and L = 64. The initial kernel will decrease in scale as L → ∞ for α L = 1 since the preactivation vectors lose variance as we discuss in Appendix E.4, resulting in slower initial training. However, we note that the final learned feature kernels are quite similar after enough training.In summary, our results indicate that the α L = 1 parameterization is the one that allows attention layers to actually be learned in the limit L → ∞, but that this parameterization leads to a less structured kernel at initialization. 
experiments in realistic settings: 
 In practice, large scale neural networks do not generally operate close to their limit. Given the costs of training large networks, one would ideally operate in a regime where there is a guarantee of consistent improvements with respect to model scale. In pursuit of this goal, we apply our theoretical findings of this paper to training language models on a larger natural language dataset, a Transformer with causal attention blocks trained on the C4 dataset [43] with Adam optimizer. As mentioned in 2.2, while our exact theoretical description of these infinite limits focus on SGD, we can implement an appropriate scaling for Adam which preserves the scale of internal feature updates. This allows us In Figure 7 (a), we sweep over each of the model dimensions independently for each parameterization of α A ∈ {1, 1 2 } on the left and right respectively. For fixed N and L, scaling H provides a similar increase in performance in both parameterization and appear to start converging to a final loss around 5, with slight benefit to α A = 1 2 . For fixed H and L, scaling N provides a similar increase in performance to scaling heads in when α A = 1, but a substantial increase when α A = 1 2 . This is in line with our predictions in Section 3.2 about the benefits of diversity across attention heads. Next, for fixed N and H, scaling L provides little to no benefit in either parameterization as predicted in Section 3.4. Finally, we inspect the sample and spatial residual stream kernels of these models before and after training and find that the kernels are identical for both α A , except for a slight difference for large N . Furthermore, they are extremely similar for large N and large H.Taken together, these results suggest that scaling different model dimensions do indeed have different effects on training dynamics and final performance. This provides groundwork for future large-scale experiments systematically investigating their trade-offs, thereby identifying compute-optimal scaling of realistic architectures in parameterizations with well-defined limits. 
discussion: 
 This paper provided analysis of the infinite head, depth and key/query dimension limits of transformer training in the feature learning regime. We showed that feature learning in µP multi-head transformers in the limit of N → ∞ collapses to single-head self-attention. At finite N and infinite heads H → ∞ we showed that there is an alternative limit which maintains a distribution over attention heads. We discussed two different large depth limits of transformer training that reduce to differential equations in the residual layer time τ . The depth scaling that maintains feature learning within all MHSA blocks (α L = 1) causes the initial kernel to lose structure from the initialization as L → ∞, but allows learning of the self-attention variables, whereas the depth scaling that preserves structure from initialization (α L = 12 ) leads to static layers.Limitations and Future Directions Currently exact theoretical analysis of the limit is focused on SGD (and can be easily extended to SGD+momentum [15] ) while Adam is currently only reasoned with rough scaling arguments rather than an exact theoretical description of the limit. Since Adam is most commonly used to train transformers, a theory of the limiting dynamics of Adam in Transformers would be an important future extension. In addition, while we provide an exact asymptotic description of network training, the limiting equations are compute intensive for realistic settings which is why we focus our empirical investigations on training large width networks in the appropriate parameterizations. Lastly our techniques assume that the number of training steps is fixed as the scaling parameters of interest (N, H, L) are taken to infinity. However, it would be important to understand learning dynamics in the regime where model size and training times are chosen to balance a compute optimal tradeoff (or perhaps even training longer than compute optimal) [8, 39, 44] . In this regime, harmful finite model-size effects become significant and comparable to the finite training horizon [10, 16, 17, 39] . Thus stress testing the ideas in this work at larger scales and longer training runs would be an important future direction of research into scaling transformer models.A Additional Figures . Initial H L (x, x ) = 16 Final H L (x, x ) = 32 = 64 = 128 Initial H L (x, x ) N = 16 Final H L (x, x ) N = 32 N = 64 N = 128 Initial H L (x, x ) L = 4 Final H L (x, x ) L = 8 L = 16 L = 32We provide an example FLAX implementation of the vision transformer and causal language model.We start by defining a fixed layernorm operation q , k = jnp . split ( qk , 2 , axis = -1) # gives q , k each of shape ( batch , heads , loc , d )v = self . v_layer ( inputs ) / jnp . sqrt ( inputs . shape [ -1] ) v = rearrange (v , 'b l ( h d ) -> b h l d ' , h = self . heads ) A = self . dim **( -self . scale_exp ) * jnp . einsum ( ' ijkl , ijml -> ijkm ' , q , k ) # batch x heads x loc x loc sigma_A = softmax ( A , axis = -1 ) out = jnp . einsum ( ' ijkl , ijlm -> ijkm ' , sigma_A , v ) # ( batch , head , loc , d ) out = rearrange ( out , 'b h l d -> b l ( h d ) ') out = self . out_layer ( out ) / jnp . sqrt ( out . shape [ -1] ) return outThe two layer MLP block is implemented as the following with ϕ = gelu nonlinearity. We also allow for a trainable positional encoding matrix. h = LN_Fixed () ( x ) h = Attention ( dim = self . dim , scale_exp = self . scale_exp , heads = self . heads ) ( h ) x = x + self . beta / self . L * h h = LN_Fixed () ( x ) h = MLP_Block ( features = self . features ) ( h ) x = x + self . beta / self . L * h return xOur vision transformer model consists of an embedding layer which is applied to each patch, a positional encoding layer, L residual layers each containing a MHSA and MLP block, a spatial pooling operation, and a readout. In this section, we heuristically work out the simple scaling analysis to justify the set of parameterizations and learning rates we consider. More detailed theoretical analysis for the limit of SGD training is provided in Appendix E where we exactly characterize the N → ∞, H → ∞ and L → ∞ limits. We consider taking heads H, inner dimension N and depth L to infinity separately and attempt to control the scale of gradients and updates.We show that the correct learning rate scaling for SGD is η = η 0 N HL 2α L -1 . For Adam, the learning rate should be scaled asη = η 0 N -1/2 H -1/2 L -1+α L . Optimizer Bulk Parameters LR First Layer Rescale Factor SGD η 0 N HL 2α L -1 L -1 2 -α Adam η 0 N -1/2 H -1/2 L -1+α L L 1-α L Table 2:The learning rates which should be 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.In this section we consider performing a single update on a single example to all weight matrices.EQUATION)where g ℓ+1 ∈ R N H and v ℓ h ∈ R N have Θ(1) entries. Thus, computing a perturbation to the forward pass we findEQUATIONThe term in the brackets is Θ(1) and we see that the perturbation from each layer contributes Θ(L -1 ).As there are L layers, this will give a total change to the final layer h L that is Θ(1).For the Attention variables, we note that theEQUATIONwhere q h , k ℓ h ∈ R N are the query and key for head h and h ∈ R N H is the residual stream preactivation. We can thus compute the changes to the keys and queries due to changes in their associated weightsEQUATION)whereH ℓ = 1 N H h ℓ • h ℓ ∼ Θ(1). Combining these changes , we find the following update to the pre-Attention variablesEQUATIONsince 1 N k • k, 1 N q • q ∼ Θ(1). This update to the attention variable due to changes in W ℓ K , W ℓ Q will clearly die out as L → ∞ unless α L = 1.For Adam, the gradient of each individual parameter entry is approximately normalized by its scale [45] . Thus the learning rate η sets the size of the updates. This is why we scale the learning rate asη = 1 L 1-α L √N H which gives the same scale updates to the weights as SGDEQUATION)Again computing the correction to the forward pass we findEQUATIONSimilarly our update generates the same scale of weight updates to the key and query weight matricesEQUATIONWe can therefore follow the identical argument to identify the scale of the change to the pre-attention variables δA ℓ h = Θ(L 1-α L ).Any parameterization with α N ∈ [ 1 2 , 1] will cause all updates to A ℓ h and entries of h ℓ+1 to be Θ N,H,L (1) across finite N . The entries of q and k only move by Θ N (1) if α A = 1 (µP scaling). However, we argue that this criterion is not strictly necessary. Rather, feature learning could alternatively be defined in terms of evolution of macroscopic variables (H, A, f , etc) rather than preactivation or key/query vector entries themselves. Table 3 summarizes two example values of α A which are of special interest for their N → ∞ limits.A(0) Update to A Update to k, q Entries α A = 1 (µP) Θ(N -1 ) Θ(1) Θ(1) α A = 1 2 Θ(1)Θ( 1) Θ(N -1 2 ) Table 3 : Two interesting choices of scaling for the attention layer exponent α A which give approximately constant updates to the attention matrices A h . The µP scaling α A = 1 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 α A = 1 2 causes the A variables to be random but still non-negligibly updated under training.The choice α A = 1 2 allows the variance of A ℓ h to be constant size as a function of N while also enabling learning of these variables. We verify these scalings in Figure 13 . Dynamical mean field theory is a method that was developed in the physics of spin glasses for dealing with dynamical systems that depend on a fixed source of disorder. The disorder could be random couplings between sites in a spin glass model [34] , random connections between neurons in a random recurrent neural network [35] , random data drawn from a distribution [37, 39] or the random initial weights in a deep neural network [15, 11] . In our case, we are interested in the last example, where the feature learning dynamics of a randomly initialized transformer is a function of the initial weights in each layer. In what follows, we will give a primer on the main objects which typically appear in a DMFT analysis (the correlation and response functions) to illustrate the main ideas of the approach.In this section, we discuss and derive the DMFT equations for the simplest possible example, a linear dynamical system with a Gaussian symmetric coupling matrix.In this example we show that the DMFT path integral is computing something non-trivial about the kinds of dynamics induced by a linear dynamical system with a random matrix. In this linear example, the DMFT path integral encodes spectral properties of the random matrix.Let's consider the simplest possible example:d dt h i (t) = 1 √ N N j=1 W ij h j (t)where W ij = W ji is a Gaussian symmetric matrix (GOE). This matrix is fixed while the state h(t) ∈ R N evolves. The path integral appraoch would tell you that in the N → ∞ limit, every neuron i has identical statistics given by the stochastic integro-differential equationEQUATIONwhere ⟨•⟩ denotes an average over the random variables u(t). In this picture, the averages ⟨⟩ over the noise can also be interpreted as averages over all N neurons in the system, each of which are independent. This stochastic equation can be used to close the evolution equations for the correlation C(t, s) and linear response function R(t, s).A generic result of this path integral DMFT picture is 1. All neurons (all variables h i ) decouple statistically. The presence of all other neurons only enters through "macroscopic" quantities C(t, s) and R(t, s) known as the correlation and response functions. The distribution of these functions over random realizations satisfies a large deviations principle where the distribution over C, R has the form p(C, R) ∼ e -N S(C,R) where S is the DMFT action obtained from the path integral method.2. Extra memory terms like t 0 R(t, s)h(s) appear which depend on the state at earlier times s < t. The Markovian (deterministic) system for p(h|W ) becomes stochastic and nonmarkovian after marginalizing p(h) = dW p(h|W )p(W ). I would argue these memory terms are not obvious apriori but are systematic to compute in this framework.Since this toy example is a linear dynamical system, one can also identify a connection between the DMFT correlation and response and spectral properties of the random matrix W . We note that the response has the formEQUATIONwhere ρ(λ) is the eigenvalue density of W . In fact a Fourier transform of our DMFT equation recovers the semicircle law ρ(λ)= 1 π ImR(iλ) = 1 2π [4 -λ 2 ] + for the eigenvalues.In general, one can think of DMFT as a more powerful version of this method that can also handle nonlinearities.In this section I will try showing how this DMFT approach can give useful insights into reasoning about learning updates which are not obvious apriori. While our paper advocates for taking depth L → ∞ in a residual network, we first thought about simply scaling depth in a standard MLP. Below we show how the proliferation of response terms gives a different predicted scaling with L than if we naively disregarded response terms.Consider a non-residual linear MLP network with µP/mean-field scaling with L hidden layers with N → ∞. Train the model for a single step of gradient descent with learning rate η on a data point (x, y) with |x| 2 = 1 and y = 1 and output multiplier 1/γ 0 . The forward pass variables h ℓ (t) and the backward pass variables g ℓ (t) are defined recursively asEQUATIONEQUATIONNow, naively, one may think that 1 √ N W ℓ (0)h ℓ (t) has entries that are independently Gaussian with covariance H ℓ (t, t ′ ). However, this is incorrect and the DMFT response functions give an additional correction.The DMFT Equations for this Model Following the approach of [15, 11] , we find the following DMFT equations for the preactivations h ℓ after 1 step of trainingEQUATIONwhere the random variables u ℓ (0), u ℓ (1) and r ℓ (0) have the following covariance structureu ℓ (0)u ℓ (0) = H ℓ-1 (0, 0) , u ℓ (1)u ℓ (0) = H ℓ-1 (1, 0) , u ℓ (1)u ℓ (1) = H ℓ-1 (1, 1)EQUATIONand the feature kernels H ℓ (t, t ′ ), G ℓ (t, t ′ ) and response functions A ℓ (t, t ′ ) have the formEQUATIONThese recursions can be solved with the initial conditions H ℓ (0, 0) = 1 and G ℓ (0, 0) = 1 which implies that H ℓ (1, 0) = 1 so thatEQUATIONUsing this equation, we findEQUATIONThis is the DMFT prediction for the scale of the feature kernels after a step of training.Neglecting the DMFT Response Gives Incorrect Depth Scalings for MLPs However, if we had neglected the DMFT response functions and approximated the dynamics asEQUATIONThe feature variance after t = 1 step of gradient descent H ℓ = h ℓ (1) 2 after t = 1 step, the final layerEQUATION)We see that without the response terms we get a completely different scaling prediction with L! Scaled Residual Networks For 1 √ L residual block scaling (α L = 1 2 ), the response functions are still important to accurately characterize the dynamics and contribute Θ L (1) corrections to the feature learning dynamics as L → ∞. However, for the 1/L block multiplier scaling, the response functions do not contribute in the limit. These facts are not-apriori obvious but follow from the DMFT analysis (either path integral or cavity approach).In this section, we derive the limiting equations for the infinite head limit H → ∞ of training with SGD. The results can be easily extended to SGD with momentum following the methods of [11, 15] .In this section we will derive the limiting equations of motion for stochastic gradient descent in the H → ∞ limit. We start by defining the loss function which we aim to minimizeEQUATIONwhere p(x) is the data distribution of interest. We note that this can be the population loss or the empirical loss on a finite collection of points. We let ∆(x, t) ≡ -∂ℓ[f (x,t)] ∂f (x,t) represent the error signal on datapoint x. At each step of training t a batch of examples B t = {x 1 (t), ..., x |Bt| (t)} is generated and used to estimate a gradient for SGD. In what follows, we let E x∼Bt represent averages over the minibatch at time t. We emphasize here that the batches B t are assumed to be given or fixed and are not averaged over as random draws from p(x), but rather our expectation simply denotes the empirical mean over the minibatch at time tEQUATIONWe start by expressing again the forward pass equations for each layer. To make this analysis more compressed while still capturing all of the interesting aspects, we will first compute the equations of motion in the absence of MLP layers (which were analyzed in prior works, see Appendix E) and layernorm (which in the limit it will only apply a deterministic affine transformation to each of the entries of the residual stream and the backward pass gradient variables as we will show explicitly in Appendix E.6). In the absence of layernorm, our forward pass has the formEQUATIONwhere the MHSA layer isEQUATIONEQUATIONTo compute the weight dynamics we again introduce the necessary gradient fields which we previously argued have Θ(1) entriesEQUATIONWe also introduce the following intermediate quantities which are necessary to characterize the backward pass through the attention layerEQUATIONWe need to break up each of the weight matrices into their initial component and their update from SGDW ℓ Oh (t) = W ℓ Oh (0) + β 0 η 0 γ 0 L 1-α L √ N H t ′ <t E x∼B t ′ s ∆(x, t ′ )g ℓ s (x, t ′ )v ℓσ hs (x, t ′ ) ⊤ W ℓ V h (t) = W ℓ V h (0) + β 0 η 0 γ 0 L 1-α L √ N H t ′ <t E x∼B t ′ ss ′ ∆(x, t ′ )σ ℓ ss ′ gℓ Ohs (x, t ′ )h ℓ s ′ (x, t ′ ) ⊤ W ℓ Kh (t) = W ℓ Kh (0) + β 0 η 0 γ 0 L 1-α L √ N H t ′ <t E x∼B t ′ ss ′ s ′′ ∆(x, t ′ )M ℓ hss ′ (x, t ′ ) σℓ ss ′ s ′′ q ℓ hs ′′ (x, t ′ )h ℓ s (x, t ′ ) ⊤ W ℓ Qh (t) = W ℓ Qh (0) + β 0 η 0 γ 0 L 1-α L √ N H t ′ <t E x∼B t ′ ss ′ s ′′ ∆(x, t ′ )M ℓ hss ′ (x, t ′ ) σℓ ss ′ s ′′ k ℓ hs (x, t ′ )h ℓ s ′′ (x, t ′ ) ⊤We can now express the residual stream asEQUATIONwhere we introduced the fieldsχℓ Os (x, t) = 1 √ H H h=1 s ′ χ ℓ Ohs ′ (x, t)σ ℓ ss ′ (x, t) , χ ℓ Ohs (x, t) = 1 √ N W ℓ Oh (0)v ℓ hs (x, t)and the kernelEQUATIONWe see that, regardless of the choice of α L , the update to the residual stream due to feature learning will scale as 1/L which is necessary for a stable infinite depth limit [11] . To track the dynamics of the value vectors, we must also track the variablesEQUATIONWe similarly find the following for the key dynamicsEQUATIONand an analogous update equation holds for the query dynamics q q ℓ hs (x, t) = χ ℓ Qhs (x, t)EQUATIONWe see that if the residual stream is frozen so that χ K , χ Q are static, the keys and queries will only evolve in the N, L → ∞ limits if α L = α A = 1 as we argued in the main text. From these equations we can deduce the pre-attention values A ℓ hss ′ (x, t). Next, we examine the dynamics of the M ℓ hs (x, t) variables which are defined asM ℓ hss ′ (x, t) = 1 N √ H g ℓ+1 s (x, t) • o ℓ hs ′ (x, t).We can verify that M ℓ h are all Θ N,H (1) throughout training by expanding the dynamics of the attention output o ℓ hs ′ (x, t).Next, we need to work out the recursions for the backward pass variables. After these equations have been worked out, we can isolate the dependence of the full dynamics on all of the initial weights.MHSA Layer We will start by differentiating through the MHSA layerEQUATION= g ℓ+1 s (x, t) + β 0 L α L ξℓ Qs (x, t) + ξℓ Ks (x, t) + ξℓ V s (x, t) + β 2 0 η 0 γ 0 L t<t ′ E x ′ ∼t ′ ∆(x ′ , t ′ )G ℓσ Oss ′ (x, x ′ , t, t ′ )h ℓ s ′ (x ′ , t ′ ) + β 2 0 η 0 γ 0 L t<t ′ E x ′ ∼t ′ ∆(x ′ , t ′ ) K ℓM σ ss ′ (x, x ′ , t, t ′ ) + Q ℓM σ ss ′ (x, x ′ , t, t ′ ) h ℓ s ′ (x ′ , t ′ ) (44)where we introduced the variablesg ℓσ Ohs (x, t) = s ′ σ ℓ ss ′ (x, t)g ℓ Ohs ′ (x, t) k ℓM σ hs (x, t) = s ′ s ′′ M ℓ hs ′ s ′′ (x, t) σℓ hs ′ s ′′′ s (x, t)k ℓ hs ′ (x, t) q ℓM σ hs (x, t) = s ′′ s ′′′ M ℓ hss ′′ (x, t) σℓ hss ′′ s ′′′ (x, t)q ℓ hs ′′′ (x, t)and their associated kernelsEQUATIONwhere we introduced the following random fieldsEQUATIONE.1.2 Why we need α A = 1 for gradient stabilityFrom the previous equation we see that regardless of the choice of α A we need to choose the variance of W ℓ Kh (0) and W ℓ Qh (0) so that 1√ N W ℓ Kh (0)q ℓ h (x, t) has O(1)entries. This means that the entries can at most Θ(1) and not Θ(N 1-α A ) as we originally stipulated in order to obtain k and q with Θ(1) entries at initialization. Thus, with this required scaling, we must have either α A = 1 and Θ(1) variance of the weights, which leads to attention variables which are Θ(N -1/2 ) at initialization or we choose α A = 1 2 and choose k, q to have entries of scale Θ(N -1/2 ) at initialization. These both lead to the same vanishing initial condition for the pre-attention variables. It thus suffices to consider µP scaling α A = 1 to study the N → ∞ limit. We stress that this effect is not visible from a simple heuristic analysis of the forward pass variables after an update like we perform in Appendix C.To summarize the previous sections, we begin by collecting all of the stochastic fields which show up in the dynamics and depend on the initial weight matrices. These quantities all come in pairs since for each matrix we need to consider the forward and backward passes through the initial matrix.The following variables are necessary to characterize the dynamics of the N H-dimensional residual streamEQUATIONFrom these variables, we can reconstruct the entire dynamics for the h, g fields. We therefore study the moment generating functional for the above primitive random variables. LetEQUATION)We multiply by the identity to enforce the definition of these random variables in terms of the initial weights. As an example, we would have for the first random variable χ ℓ+1 s (x, t) and its corresponding pair ξ ℓ s (x, t)Attention Output Matrices In this section we integrate overEQUATION)where we introduced the response functionEQUATION)Value Matrices Next, we average over the value matrices W ℓ V h (0) which givesln E W ℓ V h (0) exp - i √ N ts dx TrW ℓ V h (0) ⊤ 1 √ H χℓ V hs (x, t)h ℓ s (x, t) ⊤ + g ℓ Ohs (x, t) ξℓ V hs (x, t) ⊤ = - 1 2 tt ′ ss ′ dxdx ′ χℓ V hs (x, t) • χℓ V hs (x, t)H ℓ ss ′ (x, x ′ , t, t ′ ) + ξℓ V hs (x, t) • ξℓ V hs ′ (x ′ , t ′ )G ℓ Ohss ′ (x, x ′ , t, t ′ ) -i tt ′ ss ′ dxdx ′ χℓ V hs (x, t) • g ℓ Ohs ′ (x ′ , t ′ )R h ℓ ,ξ ℓ V hss ′ (x, x ′ , t, t ′ ) R h ℓ ,ξ ℓ V hss ′ (x, x ′ , t, t ′ ) = - i N √ H h ℓ s (x, t) • ξℓ V hs ′ (x ′ , t ′ ) (51)Key/Query Matrices Next, we need to perform the averages involving W Kh (0) which has entries resulting inln E W ℓ Kh (0) exp - i √ N ts dx TrW ℓ Kh (0) ⊤ 1 √ H χℓ Khs (x, t)h ℓ s (x, t) ⊤ + q ℓ hs (x, t) ξℓ Khs (x, t) ⊤ = - 1 2 tt ′ ss ′ dxdx ′ χℓ Khs (x, t) • χℓ Khs (x, t)H ℓ ss ′ (x, x ′ , t, t ′ ) + ξℓ Khs (x, t) • ξℓ Khs ′ (x ′ , t ′ )Q ℓ hss ′ (x, x ′ , t, t ′ ) -i tt ′ ss ′ dxdx ′ χℓ Khs (x, t) • q ℓ hs ′ (x ′ , t ′ )R h ℓ ,ξ ℓ K hss ′ (x, x ′ , t, t ′ ) R h ℓ ,ξ ℓ K hss ′ (x, x ′ , t, t ′ ) ≡ - i N √ H h ℓ s (x, t) • ξℓ Khs ′ (x ′ , t ′ ) (52)We follow an identical procedure for the query matricesW ℓ Q (0). ln E W ℓ Qh (0) exp - i √ N ts dx TrW ℓ Qh (0) ⊤ 1 √ H χℓ Qhs (x, t)h ℓ s (x, t) ⊤ + k ℓ hs (x, t) ξℓ Qhs (x, t) ⊤ = - 1 2 tt ′ ss ′ dxdx ′ χℓ Qhs (x, t) • χℓ Qhs (x, t)H ℓ ss ′ (x, x ′ , t, t ′ ) + ξℓ Qhs (x, t) • ξℓ Qhs ′ (x ′ , t ′ )K ℓ hss ′ (x, x ′ , t, t ′ ) -i htt ′ ss ′ dxdx ′ χℓ Qhs (x, t) • k ℓ hs ′ (x ′ , t ′ )R h ℓ ,ξ ℓ Q hss ′ (x, x ′ , t, t ′ ) R h ℓ ,ξ ℓ Q hss ′ (x, x ′ , t, t ′ ) ≡ - i N √ H h ℓ s (x, t) • ξℓ Qhs ′ (x ′ , t ′ ) (53)Enforce Kernel Definitions After this step, we can introduce new resolutions of the identity for each of the kernels that appeared in the above computationEQUATIONThis is repeated for all of the response functions which involve sums over N H variables1 = dR g ℓ+1 ,χ ℓ O hss ′ (x, x ′ , t, t ′ )dR v ℓ ,ξ ℓ O s ′ s (x ′ , x, t ′ , t) 2πiN -1 exp -R v ℓ ,ξ ℓ O hs ′ s (x ′ , x, t ′ , t) N R g ℓ+1 ,ξ ℓ O hss ′ (x, x ′ , t, t ′ ) + i √ H g ℓ+1 s (x, t) • χℓ Ohs ′ (x ′ , t ′ ) 1 = dR h ℓ ξ ℓ V hss ′ (x, x ′ , t, t ′ )dR g ℓ O χ ℓ V hs ′ s (x ′ , x, t ′ , t) 2πiN -1 exp -R g ℓ O χ ℓ V hs ′ s (x ′ , x, t ′ , t) N R h ℓ ξ ℓ V hss ′ (x, x ′ , t, t ′ ) + i √ H h ℓ s (x, t) • ξℓ V hs ′ (x ′ , t ′ ) 1 = dR h ℓ ξ ℓ K hss ′ (x, x ′ , t, t ′ )dR q ℓ χ ℓ K hs ′ s (x ′ , x, t ′ , t) 2πiN -1 exp -R q ℓ χ ℓ K hs ′ s (x ′ , x, t ′ , t) N R h ℓ ξ ℓ K hss ′ (x, x ′ , t, t ′ ) + i √ H h ℓ s (x, t) • ξℓ Khs ′ (x ′ , t ′ ) 1 = dR h ℓ ξ ℓ Q hss ′ (x, x ′ , t, t ′ )dR k ℓ χ ℓ Q hs ′ s (x ′ , x, t ′ , t) 2πiN -1 exp -R k ℓ χ ℓ Q hs ′ s (x ′ , x, t ′ , t) N R h ℓ ξ ℓ Q hss ′ (x, x ′ , t, t ′ ) + i √ H h ℓ s (x, t) • ξℓ Qhs ′ (x ′ , t ′ )There are other kernels which are only relevant within a single head including{Q ℓ h , K ℓ h , A ℓ h , M ℓ h }. These 1 = dQ ℓ hss ′ (x, x ′ , t, t ′ )d Qℓ hss ′ (x, x ′ , t, t ′ ) 2πiN -1 exp Qℓ hss ′ (x, x ′ , t, t ′ ) N Q ℓ hss ′ (x, x ′ , t, t ′ ) -q ℓ hs (x, t) • q ℓ hs ′ (x ′ , t ′ ) 1 = dK ℓ hss ′ (x, x ′ , t, t ′ )d Kℓ hss ′ (x, x ′ , t, t ′ ) 2πiN -1 exp Kℓ hss ′ (x, x ′ , t, t ′ ) N K ℓ hss ′ (x, x ′ , t, t ′ ) -k ℓ hs (x, t) • k ℓ hs ′ (x ′ , t ′ ) 1 = dM ℓ hss ′ (x, t)d M ℓ hss ′ (x, t) 2πiN -1 exp M ℓ hss ′ (x, t) N M ℓ hss ′ (x, t) - 1 √ H gℓ s (x, t) • o ℓ hs ′ (x, t) 1 = dA ℓ hss ′ (x, t)d Âℓ hss ′ (x, t) 2πiN -1 exp Âℓ hss ′ (x, t) N A ℓ hss ′ (x, t) -k ℓ hs (x, t) • q ℓ hs ′ (x, t)We now combine all of the order parameters into a large collection Q which are vectorized over all layer, time, spatial and sample indicesEQUATION)After introducing this collection of order parameters, our original MGF satisfies a large deviation principle withEQUATIONwhere the DMFT action S(Q) has the formEQUATIONThe residual stream single site moment generating functional has the formZ res = ℓhsxt d χOhs (x, t)dχ Ohs (x, t) 2π d ξV hs (x, t)dξ V hs (x, t) 2π d ξKhs (x, t)dξ Khs (x, t) 2π d ξQhs (x, t)dξ Qhs (x, t) 2π exp - L ℓ=1 tt ′ ss ′ dxdx ′ Ĥℓ ss ′ (x, x ′ , t, t ′ )h ℓ s (x, t)h ℓ s ′ (x ′ , t ′ ) exp - L ℓ=1 tt ′ ss ′ dxdx ′ Ĝℓ ss ′ (x, x ′ , t, t ′ )g ℓ s (x, t)g ℓ s ′ (x ′ , t ′ ) exp   - L ℓ=1 h tss ′ dx M ℓ hss ′ (x, t)g ℓ+1 s (x, t)o ℓ hs ′ (x, t)   exp   - 1 2 ℓh tt ′ ss ′ dxdx ′ χℓ Ohs (x, t) χℓ Ohs ′ (x ′ , t ′ )V ℓ hss ′ (x, x ′ , t, t ′ )   exp   - 1 2 ℓh tt ′ ss ′ dxdx ′ ξℓ V hs (x, t) ξℓ V hs ′ (x ′ , t ′ )G ℓ Ohss ′ (x, x ′ , t, t ′ )   exp   - 1 2 ℓh tt ′ ss ′ dxdx ′ ξℓ Khs (x, t) ξℓ Khs ′ (x ′ , t ′ )Q ℓ hss ′ (x, x ′ , t, t ′ )   exp   - 1 2 ℓh tt ′ ss ′ dxdx ′ ξℓ Qhs (x, t) ξℓ Qhs ′ (x ′ , t ′ )K ℓ hss ′ (x, x ′ , t, t ′ )   exp   i ℓh ts dx χℓ Ohs (x, t)χ ℓ Ohs (x, t) + ξℓ V hs (x, t)ξ ℓ V hs (x, t)   exp   i ℓh ts dx ξℓ Qhs (x, t)ξ ℓ Qhs (x, t) + ξℓ Khs (x, t)ξ ℓ Khs (x, t)   exp   - i √ H ℓh tt ′ ss ′ dxdx ′ R v ℓ ,ξ ℓ O hss ′ (x, x ′ , t, t ′ ) χℓ Ohs (x, t)g ℓ+1 s ′ (x ′ , t ′ )   exp   - i √ H ℓh tt ′ ss ′ dxdx ′ R g ℓ O χ ℓ V hss ′ (x, x ′ , t, t ′ ) ξℓ V hs (x, t)h ℓ s ′ (x ′ , t ′ )   exp   - i √ H ℓh tt ′ ss ′ dxdx ′ R q ℓ χ ℓ K hss ′ (x, x ′ , t, t ′ ) ξℓ Khs (x, t)h ℓ s ′ (x ′ , t ′ )   exp   - i √ H ℓh tt ′ ss ′ dxdx ′ R k ℓ χ ℓ Q hss ′ (x, x ′ , t, t ′ ) ξℓ Qhs (x, t)h ℓ s ′ (x ′ , t ′ )   (58)Though this expression is cumbersome, we will show that, since Ĥ, Ĝ vanish at their saddle point this MGF merely encodes the following statistical description of the fields of interest such asEQUATIONWe note that h only depends on χℓO = 1 √ H H h=1 χ ℓ Oh σ ℓ h so it has the form χℓ O = 1 √ H H h=1 u ℓ Oh σ ℓ h + Rv ℓ ξ ℓ O g ℓ , Rv ℓ ξ ℓ O = 1 H H h=1 R v ℓ ,ξ ℓ O h σ ℓ h (60)This fact will be important when we take the large head limit with N fixed in Appendix E.3.Next, we analyze the single site MGF for the hidden MHSA layersEQUATIONE.2 Infinite N (Key/Query dimension) LimitFirst, we take the N → ∞ limit with H, L fixed. This can be obtained with a simple saddle point procedure using the action in the form written in the previous section. This calculation exactly mimics prior works [15, 11] where all of the order parameters Q take on their values at the saddle pointEQUATIONSaddle Point Values for Order Parameters Under this saddle point, all of the order parameters presented will concentrate and all neurons will become statistically independent. The governing equations for the order parameters Q ⋆ are given in terms of averages ⟨•⟩ over the single site densities defined by the moment generating functionals Z and have the formEQUATIONSingle Site Stochastic Processes Our fields of interest will obey the stochastic dynamicsEQUATIONResidual Stream The forward pass residual variables obey the following stochastic processEQUATIONwhereEQUATIONthe backward pass satisfiesEQUATIONThe keys and queries have dynamicsEQUATIONE.2.1 Multi-head Attention is Single-Head Attention as N → ∞In this section, we use the derived saddle point equations for the N → ∞ limit and argue that they imply that all heads in the MHSA layer learn identical attention matrices and contribute the same feature updates to the residual stream. To do so, we proceed in a three step inductive argument.1. First, we show that at initialization, all key, query, value and attention matrices {Q h , K h , V h , A h , M h } are equal across heads.2. Next, we show inductively that if these quantities {Q h , K h , V h , A h , M h } are identical across heads up to some time, then that implies that the response functions R h are also identical across heads up to that time.3. Lastly, we show that if the response functions R h are identical up to some time, then that implies that the MHSA kernels {Q h , K h , V h , A h , M h } will also be identical across heads at future times.First, we note that, at initialization, all of the MHSA kernels are identical across heads sinceEQUATIONNext, we need to analyze the response functions under an inductive hypothesis on the equality of the MHSA kernels. We start by noting that all response functions are causal so we can group the response functions that arise from χ ℓ Oh with the feature learning update to the residual stream, writing the following compressed equation for the forward and backward passesh ℓ+1 s (x, t) = h ℓ s (x, t) + 1 √ H H h=1 s ′ u ℓ Ohs ′ (x, t)σ ℓ hss ′ (x, t) + η 0 γ 0 β 2 0 L -1 t ′ <t s ′ dx ′ C h ℓ ss ′ (x, x ′ , t, t ′ )g ℓ+1 s ′ (x ′ , t ′ ) g ℓ s (x, t) = g ℓ s (x, t) + 1 √ H hs ′ r ℓ V hs ′ (x, t)σ ℓ hs ′ s (x, t) + 1 √ H hs ′ s ′′ r ℓ Qhs ′ (x, t) σℓ hs ′ s ′′ s (x, t)M ℓ hs ′ s ′′ (x, t) + 1 √ H hs ′′ s ′′′ r ℓ Khs ′′′ (x, t) σℓ hss ′′ s ′′′ (x, t)M ℓ hss ′′ (x, t) + η 0 γ 0 β 2 0 L -1 t ′ <t s ′ dx ′ C g ℓ ss ′ (x, x ′ , t, t ′ )h ℓ s ′ (x ′ , t ′ )where C h ℓ and C g ℓ only involve deterministic head-averaged kernels and thus do not carry a h index. We now derive useful response function identitiesEQUATIONThese equations give the needed response function R g ℓ+1 χ ℓ O . From these equations we immediately see that if σ ℓ hss ′ (x, t) = σ ℓ h ′ ss ′ (x, t) (which holds under our inductive hypothesis) then Rg ℓ+1 χ ℓ O h = R g ℓ+1 χ ℓ O h ′. This same argument is repeated for all other response functions that are computed as derivatives of residual stream variables. We have thus found that∀h, h ′ ∈ [H] , R g ℓ+1 χ ℓ O h = R g ℓ+1 χ ℓ O h ′ , R h ℓ ξ ℓ V h = R h ℓ ξ ℓ V h ′ , R h ℓ ξ ℓ K h = R h ℓ ξ ℓ K h ′ , R h ℓ ξ ℓ Q h = R h ℓ ξ ℓ Q h ′Now, we can analyze the dynamics of the keys, queries and values within a head using the above property and the original inductive hypothesis that A h = A h ′ , M h = M h ′ ... for times less than t. We will now prove that this implies that these variables will remain the same. We start by examining the keys and queries which have the formEQUATIONEQUATIONwhere C k ℓ ss ′ (x, x ′ , t, t ′ ) and C q ℓ ss ′ (x, x ′ , t, t ′ ) are two operators that do not carry a h (are the same across all heads). These relations can be viewed as a linear system of equations. We let k ℓ h = Vec{k ℓ hs (x, t)} sxt and analogously C k ℓ = Mat{C k ℓ ss ′ (x, x ′ , t, t ′ )} as in the calculations of Bordelon and Pehlevan [15, 41] . Using this shorthand, we can express the key/query and attention kernels asEQUATIONWe thus see that the final kernels K ℓ h , Q ℓ h , A ℓ h are all identical across heads. An identical argument can be carried out for the value kernel V ℓ h and the M ℓ h order parameter.In this section, we compute the infinite head limit with N, L fixed. This limit is more technically involved than the N → ∞ limit which required only a simple saddle point of the full DMFT action over all kernels. At finite N we cannot use this technique since the kernels within MHSA blocks are random variables. However, as was shown in Bordelon and Pehlevan [17] , the DMFT action still contains the necessary information to characterize the distribution over order parameters at finite N .In the case of transformers with infinitely many heads, a subset of the order parameters introduced in the previous section will still concentrate as H → ∞ includingH ℓ ss ′ (x, x ′ , t, t ′ ) = 1 N H h ℓ s (x, t) • h ℓ s (x, t) G ℓ ss ′ (x, x ′ , t, t ′ ) = 1 N H g ℓ s (x, t) • g ℓ s (x, t) V ℓσ ss ′ (x, x ′ , t, t ′ ) = 1 H H h=1 s ′′ s ′′′ V ℓσ hs ′′ s ′′′ (x, x ′ , t, t ′ )σ ℓ hss ′′ (x, t)σ ℓ hss ′′′ (x ′ , t ′ )and many more correlation and response functions. We will call the full collection of all the necessary head-averaged order parameters Q global . Further, not all of the stochastic fields will be relevant to characterize the residual stream. Specifically, only head-averaged fields χℓ O , ξℓ V , ξℓ K , ξℓ Q are relevant. For example, the first of these is defined asEQUATIONWe note that the above equation is true at any value of N but the covariance of u ℓ h and the response functions Rv ℓ ξ ℓ O hare random variables [17] . However, the residual stream only depends upon collective, head-averaged variablesEQUATIONEQUATIONThe intuition behind the large H limit is that, even though σ ℓ h and Rv ℓ ξ ℓ O hare random variables, there should be a central limit theorem for ūℓ O and a law of large numbers for Rv ℓ ξ ℓ O σ ss ′ (x, x ′ , t, t ′ ).Based on the intuition developed in the previous section, we now derive an alternative DMFT action by tracking the moment generating functional for the head-averaged random fields that occur on the residual stream χℓ O , ξℓ V , ξℓ K , ξℓ Q . To characterize this joint distribution, we must also of course keep track of the random fields within the MHSA blocks such as {χ ℓ Qh , χ ℓ Kh , χ ℓ V h , ξ ℓ Oh } h∈ [H] . Repeating the path integral setup of the previous section, we need to performing averages over all initial weights such asEQUATIONIt is clear that from this integral that the relevant self-averaging order parameters areEQUATIONWe repeat this same procedure for all collections of weights and arrive at the following set of order parametersEQUATIONWe expect that these order parameters will be self-averaging since they involve averages over H variables. However, the other variables we introduced {A h , Q h , K h , V h } will not concentrate at finite N and will instead behave as random variables.After introducing these order parameters we find that the moment generating functional has the formEQUATIONwhere S has the formEQUATION)The single site moment generating function for the residual stream has the form Zd ξℓ V s (x, t)d ξℓ V s (x, t) 2π exp - ℓtt ′ ss ′ dxdx ′ h ℓ s (x, t)h ℓ s ′ (x ′ , t ′ ) Ĥℓ ss ′ (x, x ′ , t, t ′ ) + g ℓ s (x, t)g ℓ s ′ (x ′ , t ′ ) Ĝℓ ss ′ (x, x ′ , t, t ′ ) exp - 1 2 ℓtt ′ ss ′ dxdx ′ χℓ Os (x, t) χℓ Os ′ (x ′ , t ′ )V ℓσ ss ′ (x, x ′ , t, t ′ ) + ξℓ V s (x, t) ξℓ V s ′ (x ′ , t ′ )G ℓσ Oss ′ (x, x ′ , t, t ′ ) exp - 1 2 ℓtt ′ ss ′ dxdx ′ ξℓ Ks (x, t) ξℓ Ks ′ (x ′ , t ′ )Q ℓM σ ss ′ (x, x ′ , t, t ′ ) + ξℓ Qs (x, t) ξℓ Qs ′ (x ′ , t ′ )K ℓM σ ss ′ (x, x ′ , t, t ′ ) exp -i ℓtt ′ ss ′ dxdx ′ χℓ Os (x, t)g ℓ+1 s ′ (x ′ , t ′ ) Rv ℓ ξ ℓ O ss ′ (x, x ′ , t, t ′ ) + ξℓ V s (x, t)h ℓ s ′ (x ′ , t ′ ) Rg ℓ O χ ℓ V ss ′ (x, x ′ , t, t ′ ) exp -i ℓtt ′ ss ′ dxdx ′ ξℓ Ks (x, t)h ℓ+1 s ′ (x ′ , t ′ ) Rq ℓ χ ℓ K ss ′ (x, x ′ , t, t ′ ) + ξℓ Qs (x, t)h ℓ s ′ (x ′ , t ′ ) Rk ℓ χ ℓ Q ss ′ (x, x ′ ,We can express the MHSA single-head partition functions Z MHSA in terms of the remaining order parameters within each head that will no longer concentrate at finite NEQUATIONAfter introducing these order parameters, we have Z MHSA = dQ ℓ MHSA exp N S MHSA (Q ℓ MHSA ) S MHSA = tt ′ ss ′ dxdx ′ [Q ℓ ss ′ (x, x ′ , t, t ′ ) Qℓ ss ′ (x, x ′ , t, t ′ ) + K ℓ ss ′ (x, x ′ , t, t ′ ) Kℓ ss ′ (x, x ′ , t, t ′ )] + tt ′ ss ′ dxdx ′ [V ℓ ss ′ (x, x ′ , t, t ′ ) V ℓ ss ′ (x, x ′ , t, t ′ ) + G ℓ Oss ′ (x, x ′ ,EQUATIONEQUATIONThe saddle point equations for this limit are computed as derivatives with respect to Q global only, reflecting that head-averages will converge as H → ∞EQUATIONThe final saddle point equations are given in terms of averages over the distribution of heads defined by Z MHSA which we denote as ⟨•⟩ MHSA as well as averages over the residual stream which we denote as ⟨•⟩.These equations give the following (we suppress the sequence indices to simplify the final expressions)EQUATION)Residual Stream Dynamics The residual stream satisfies the following single-site dynamicsEQUATIONThis matches the result provided in the main text which introduces a compressedEQUATIONwhere p t (x) = 1 |Bt| x ′ ∈Bt δ(xx ′ ) denotes the uniform distribution over the batch B t .In this section, we discuss the two large L limits. This can be derived formally in two distinct ways. First, one could start with the initial For this section, it suffices to reason about the scale of the Gaussian noise which appears in the residual stream and the contribution from the response functions.Deriving the infinite depth limits To gain intuition for the large H, L → ∞ limit, we use the fact that the random variables χ ℓ = u ℓ + R ℓ g ℓ decompose into a Gaussian u ℓ which are uncorrelated across layers and a linear response R ℓ are response functions. This implies thatEQUATIONWe first note that the sum of the Gaussians is a zero-mean random variable with standard deviationEQUATIONThus, this integrated random variable will vanish unless α L = 1 2 . Next, we can investigate the scale of the residual stream response functions. For instance∂h ℓ ∂u k = O L -α L , ∂h ℓ ∂r k = O L -α L , ∂g ℓ ∂u k = O L -α L , ∂g ℓ ∂r k = O L -α L 1 L α L ℓ k=1 R k g k = O L 1-2α L (92)As a consequence, we see that the effect of the Gaussian and linear response terms will vanish as L → ∞ provided that α L > 1 2 . We will consider first, the case where α = 1 which gives an ODE like limit for the residual updates before moving onto the more involved α L = 1 2 case. To formally take the L → ∞ limit, we redefine all of the preactivation fields and kernels in terms of layer time τ defined asEQUATIONFor example, the residual kernels are defined asEQUATIONThe finite difference equations for the residual updates L(h ℓ+1h ℓ ) ∼ ∂ ∂τ h(τ ) become differential updates (either SDE-like or ODE-like depending on α L ) [46, 11] .E.4.2 ODE Limit α L = 1First, we investigate the case of α L = 1. In this case, the H, L → ∞ limit results in a complete disappearance of the χℓO , ξℓ V , ξℓ K , ξℓ Q fields. ∂ τ h s (τ, x, t) = η 0 γ 0 β 2 0 t ′ <t E x ′ ∼Bt ∆(x ′ , t ′ )V σ ss ′ (τ, x, x ′ , t, t ′ )g s ′ (τ ′ , x ′ , t ′ ) -∂ τ g s (τ, x, t) = η 0 γ 0 β 2 0 t ′ <t E x ′ ∼Bt ∆(x ′ , t ′ )[G σ ss ′ (τ, x, x ′ , t, t ′ ) + K M σ ss ′ (τ, x, x ′ , t, t ′ )]g s ′ (τ ′ , x ′ , t ′ ) + η 0 γ 0 β 2 0 t ′ <t E x ′ ∼Bt ∆(x ′ , t ′ )[Q M σ ss ′ (τ, x, x ′ , t, t ′ )]g s ′ (τ ′ , x ′ , t ′ ) (95) E.4.3 SDE Limit α L = 1 2This α L = 1 2 limit is more technically involved since neither the Gaussian terms from the DMFT nor the response functions vanish. For the Gaussian terms, we note that the sums of the independent Gaussians are all multiplied by 1 √ L ∼ √ dτ , which can be interpreted as integrated Brownian motion in the limitEQUATIONFollowing the derivation of Bordelon et al. [11] , which maintains the exact dependence on the full integrated response and provides the result as an integrated SDEEQUATIONCombining the forward pass equations from the previous two subsection recovers the Result 3 of the main text. This is combined with a complementary equation for the backward pass.Adding the MLP block to the residual stream can also be easily handled with the methods of the preceeding sections. The forward pass equations in this case take the formEQUATIONwhere the MLP layer isEQUATIONThe following gradient fields are necessaryEQUATIONEQUATIONThe MLP hidden layer dynamics is much simpler to characterize and resembles the structure analyzed in prior works on infinite width networks [15] .EQUATIONAgain, we see that the inner dynamics for hℓ,1 due to the weight updates in this layer scale as L -1+α L , suggesting the need to choose α L = 1 if we desire this hidden layer to contribute to the representational updates.MLP Layer Gradients For the MLP layer we have the simpler backpropagation equationsEQUATIONThe components of these fields that depend on initial conditions areEQUATIONMLP Matrices After utilizing these resolutions of the identity for all s, x, t, we can integrate over the weights W ℓ,2 (0)EQUATIONwhere we introduced the response functionEQUATIONWe can perform an identical step to integrate over W ℓ,1 (0). This gives usEQUATIONwhere we introduced acts as the following in the large H limitEQUATION∂ h ∂h ⊤ = 1 √ σ 2 + ϵ I - 1 N H 11 ⊤ - 1 N H 1 (σ 2 + ϵ) 3/2 [h -µ1] [h -µ1] ⊤ (110)In the limit of H → ∞ the variables µ = 1 N H h • 1 and σ 2 = 1 N H |h -µ1| 2 will become deterministic over random initializations. We thus just have to consider how these types of vectors act on gradientsEQUATIONEach of these operations will lead to one inner product that will be self averagingµ g = 1 N H 1 • g or 1 N H [h -µ1]⊤ g . Thus this operation will not alter the backward pass in terms of scaling.Each of the experimental runs performed in this paper were all performed on single NVIDIA H100 GPU. Each run of the full CIFAR-5M took anywhere from 5 minutes to 1 hour depending on model size. Each run of the C4 training took anywhere from 1 hour to 6 hours depending on model size and total amount of training steps.The language model trained on C4 used a context length of 256 and the GPT-tokenizer from Huggingface. Sequences that were too short were concatenated with other sentences to reach the full context length rather than padding the end of the sequence. We use trainable positional encodings and separate embedding and decoding parameters as implemented in Appendix F.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We provide code in the uploaded supplementary material. Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We describe the number of ensembles and general experimental details (SGD vs Adam, depth, width, head number etc) in all of our experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We provide errorbars wherever appropriate, usually averaging and measuring standard deviation over different initializations.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide a section mentioning the compute resources used. All experiments can be performed on a single Nvidia A100 or H100 GPU. The details on run times for each experiment are provided in Appendix F.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We are conforming to the code of ethics and preserve our anonymity.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This work is of a theoretical nature and is about general scientific understanding of transformer models.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?38th Conference on Neural Information Processing Systems (NeurIPS 2024).Training time, sample size, sequence length/spatial dimension are all treated as fixed Θ(1) quantities. 
 
 Appendix:  


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

Reviews: review: paper_summary: paper_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:
- taking the embedding dimension to infinity can make heads redundant with each other, while fixing embedding dimension and scaling head count avoids that issue)
- 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

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

I'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.
strength_weakness: strength_weakness: 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...

### **Technical tools**

The 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!

Now 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?

### **Exposition and clarity**

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

### **Missing related work on non-asymptotic approaches to scaling**

The 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.
questions: questions: Please see the Weaknesses section.
limitations: 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".

By the way, I am willing to upgrade my score. But will be keen to see and engage with the opinions of the other reviewers.
paper_summary: paper_summary: The authors identify parameterizations that lead to nontrivial feature learning as the limits of $N, H, L \to \infty$. Specifically, the study demonstrates the following:

- Under the limit $N\to\infty$ the $\mu P$ rule is required, causing all heads to collapse into the same dynamics.
- Under the limit $H\to\infty$, each head becomes statistically independent.
- The scaling under the depth limit $L\to\infty$ is also analyzed.
strength_weakness: strength_weakness: - It is unclear whether considering the limit of the number of heads $H$ is practically relevant, as $H$ is typically set to around or fewer than 100 in real-world applications. This raises questions about whether the theory accurately explains the behavior of large language models (LLMs).
- The following paper is also relevant:

  - Lénaïc Chizat, Praneeth Netrapalli. The Feature Speed Formula: A flexible approach to scale hyper-parameters of deep neural networks. 2024.
questions: questions: - Line 78: What is $\gamma_0$? Could you elaborate on this further?
- What is the maximum model size (i.e., number of parameter) used in the experiments?
- Do the results hold for any time t? If so, this feature is worth emphasizing, as most papers, such as [L. Chizat and P. Netrapalli (2024)] and [G. Yang and E.J. Hu (2021)], focus on the initialization phase.
limitations: limitations: Limitations are addressed in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The paper is hard to understand, the notations are convoluted and most of the community might not be familiar with dynamical mean field theory , the authors might want to offer more background in the main text
questions: questions: 1. In line 78, it is mentioned that gamma_0 controls the rate of feature learning, but I couldn't find it in equation (1)
2. Why according to Table 1, SGD's learning rate scales with N and H, but Adam's learning rate scales down when N and H decreases, any intuition to understand this phenomenon?
3. Why 3/2 appears in eq3
4. Why the limit can be straightforwardly computed from the saddle point of DMFT action?
limitations: limitations: Limitations are addressed in the last session of the paper
paper_summary: paper_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.
strength_weakness: strength_weakness: See Questions.
questions: 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?
* 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?
* 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?
* What are some previous results or potential approaches towards analyzing limiting dynamics of neural networks under Adam?
limitations: limitations: A Limitations section is provided in Section 5.

Comments: title: Some non-trivial DMFT conclusions
Some Non-Vacuous Path Integral Examples
comment: Thank you for your comments! We think including the alternative cavity derivation of our DMFT will be worth including as this approach (in other settings at least) can be made rigorous, such as in the work https://arxiv.org/abs/2210.06591. 

We will also aim to improve the presentation where possible and motivate our calculations more clearly.
Hi---I'm sorry for the delay.

I think this is a good piece of science, and the principles you raise are very very interesting. I'm not sure how rigorous your calculations are (again, I haven't checked them) but I certainly think this paper is worth presenting on a scientific level. I want to share some thoughts on your work that are intended to be well-meaning although perhaps provocative:

- I talked to a mathematician friend about path integrals, and he told me mathematicians view them as non-rigorous and don't do them. He told they worked out alternate ways to do the same calculations using Morse theory or Floer homology. I don't know to what extent this is academic tribalism or whether there's something to it, but I thought I'd share the perspective.

- I think that if you want the wider community to seriously engage with your methods, you might need to try to find a killer app still and also put a lot of effort into honing the presentation. I say this since I do have the sense that there are other ways to do the scaling calculations that you're doing, and perhaps simpler ones. Of course I may be wrong about this, and again I do like the science that you're doing.

Regarding your rebuttal, I think the changes that you've committed to will strengthen the paper, and I'm comfortable increasing my score.
Hi---sorry, I've been engaging a lot with 4 other papers in my batch, and it's prevented me from yet reaching yours. I plan to respond thoroughly, so stay tuned, but for now consider this acknowledgement of your rebuttal.
As the discussion period is ending soon, we were hoping to follow up to see if our response answered the reviewer's questions. If so, we would hope that they would be willing to increase their score. If not, we would be happy to address any additional questions or concerns. We thank you again for your time and reviews.
As the discussion period is ending soon, we were hoping to follow up to see if our response answered the reviewer's questions. If so, we would hope that they would be willing to increase their score. If not, we would be happy to address any additional questions or concerns. We thank you again for your time and reviews.
Thank you for the detailed explanations including interesting examples and new (to me) ideas. I will maintain my positive review of the paper.
Thank you for your response! I will keep my already positive score
### A very simple example: GOE/Wigner Linear Dynamics

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

Let's consider the simplest possible example: $\frac{d}{dt} h_i(t) = \frac{1}{\sqrt N} \sum_{j=1}^N W_{ij} h_j(t)$ where $W_{ij} = W_{ji}$ is a Gaussian symmetric matrix (GOE). This matrix is **fixed** while the state $h(t) \in \mathbb{R}^N$ evolves. The path integral approch would tell you that in the $N \to \infty$ limit, every neuron $i$ has identical statistics given by the stochastic integro-differential equation
\begin{align}
    &\partial_t h(t) = u(t) + \int_0^t ds R(t,s) h(s)  \ , \ u(t) \sim \mathcal{GP}(0, C(t,s))
    \\
    &C(t,s) = \left< h(t) h(s) \right> \ , \ R(t,s) = \left< \frac{\delta h(t)}{\delta u(s)} \right>
\end{align}
where $\left< \cdot \right>$ denotes an average over the random variables $u(t)$. This stochastic equation can be used to close the evolution equations for the correlation $C(t,s)$ and linear response function $R(t,s)$. 

A generic result of this path integral DMFT picture is
1. All neurons decouple statistically. The presence of all other neurons only enters through "macroscopic" quantities $C(t,s)$ and $R(t,s)$ known as the correlation and response functions. The distribution of these functions over random realizations satisfies a large deviations principle $p(C,R) \sim e^{- N S(C,R)}$ where $S$ is the DMFT action. 
2. Extra *memory terms* like $\int_0^t R(t,s) h(s)$ appear which depend on the state at earlier times $s < t$. The Markovian (deterministic) system for $p(h|W)$ becomes stochastic and non-markovian after marginalizing $p(h) = \int dW p(h|W) p(W)$. I would argue these memory terms are not obvious apriori but are systematic to compute in this framework.

Since this toy example is a **linear dynamical system**, one could also obtain the correlation $C(t,s)$ and response $R(t,s)= \frac{1}{N} \text{Tr} \exp\left( W (t-s) \right) = \int d\lambda \rho(\lambda) e^{ \lambda (t-s)}$ where $\rho(\lambda)$ is the eigenvalue density of $W$. In fact a Fourier transform of our DMFT equation recovers the semicircle law $\rho(\lambda) = \frac{1}{\pi} \text{Im} R(i \lambda) = \frac{1}{2\pi} \sqrt{[4-\lambda^2]_+}$ for the eigenvalues. 

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

### Do Memory / Response Terms Matter in Deep Networks?

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

Consider a non-residual linear MLP network with $\mu$P/mean-field scaling with $L$ hidden layers with $N \to \infty$. Train the model for a single step of gradient descent with learning rate $\eta$ on a data point $(x,y)$ with $|x|^2 = 1$ and $y=1$. The feature variance after $t=1$ step of gradient descent $H^{\ell} = \left< h^\ell(1)^2 \right>$ after $t=1$ step, the final layer 
\begin{align}
    H^{L} \sim \begin{cases} 1 +  \frac{1}{3} \eta^2 \gamma_0^2 \ L^3 & \text{DMFT Response Included (Full DMFT)}
    \\
    1 + \eta^2 \gamma_0^2 \ L & \text{DMFT Response Neglected }
    \end{cases}
\end{align}

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


For $\frac{1}{\sqrt L}$ residual block scaling, the response functions are still very important and contribute $\Theta_L(1)$ corrections to the feature learning dynamics as $L \to \infty$.  However, for the $1/L$ block multiplier scaling, the response functions do not contribute in the limit. These facts are **not apriori obvious** (to us at least) but follow from the DMFT analysis (either path integral or cavity approach).
### A very simple example: GOE/Wigner Linear Dynamics

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

Let's consider the simplest possible example: $\frac{d}{dt} h_i(t) = \frac{1}{\sqrt N} \sum_{j=1}^N W_{ij} h_j(t)$ where $W_{ij} = W_{ji}$ is a Gaussian symmetric matrix (GOE). This matrix is **fixed** while the state $h(t) \in \mathbb{R}^N$ evolves. The path integral approch would tell you that in the $N \to \infty$ limit, every neuron $i$ has identical statistics given by the stochastic integro-differential equation
\begin{align}
    &\partial_t h(t) = u(t) + \int_0^t ds R(t,s) h(s)  \ , \ u(t) \sim \mathcal{GP}(0, C(t,s))
    \\
    &C(t,s) = \left< h(t) h(s) \right> \ , \ R(t,s) = \left< \frac{\delta h(t)}{\delta u(s)} \right>
\end{align}
where $\left< \cdot \right>$ denotes an average over the random variables $u(t)$. This stochastic equation can be used to close the evolution equations for the correlation $C(t,s)$ and linear response function $R(t,s)$. 

A generic result of this path integral DMFT picture is
1. All neurons decouple statistically. The presence of all other neurons only enters through "macroscopic" quantities $C(t,s)$ and $R(t,s)$ known as the correlation and response functions. The distribution of these functions over random realizations satisfies a large deviations principle $p(C,R) \sim e^{- N S(C,R)}$ where $S$ is the DMFT action. 
2. Extra *memory terms* like $\int_0^t R(t,s) h(s)$ appear which depend on the state at earlier times $s < t$. The Markovian (deterministic) system for $p(h|W)$ becomes stochastic and non-markovian after marginalizing $p(h) = \int dW p(h|W) p(W)$. I would argue these memory terms are not obvious apriori but are systematic to compute in this framework.

Since this toy example is a **linear dynamical system**, one could also obtain the correlation $C(t,s)$ and response $R(t,s)= \frac{1}{N} \text{Tr} \exp\left( W (t-s) \right) = \int d\lambda \rho(\lambda) e^{ \lambda (t-s)}$ where $\rho(\lambda)$ is the eigenvalue density of $W$. In fact a Fourier transform of our DMFT equation recovers the semicircle law $\rho(\lambda) = \frac{1}{\pi} \text{Im} R(i \lambda) = \frac{1}{2\pi} \sqrt{[4-\lambda^2]_+}$ for the eigenvalues. 

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

### Do Memory / Response Terms Matter in Deep Networks?

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

Consider a non-residual linear MLP network with $\mu$P/mean-field scaling with $L$ hidden layers with $N \to \infty$. Train the model for a single step of gradient descent with learning rate $\eta$ on a data point $(x,y)$ with $|x|^2 = 1$ and $y=1$. The feature variance after $t=1$ step of gradient descent $H^{\ell} = \left< h^\ell(1)^2 \right>$ after $t=1$ step, the final layer 
\begin{align}
    H^{L} \sim \begin{cases} 1 +  \frac{1}{3} \eta^2 \gamma_0^2 \ L^3 & \text{DMFT Response Included (Full DMFT)}
    \\
    1 + \eta^2 \gamma_0^2 \ L & \text{DMFT Response Neglected }
    \end{cases}
\end{align}

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


For $\frac{1}{\sqrt L}$ residual block scaling, the response functions are still very important and contribute $\Theta_L(1)$ corrections to the feature learning dynamics as $L \to \infty$.  However, for the $1/L$ block multiplier scaling, the response functions do not contribute in the limit. These facts are **not apriori obvious** (to us at least) but follow from the DMFT analysis (either path integral or cavity approach).