# Infinite Limits of Multi-head Transformer Dynamics

Blake Bordelon, Hamza Chaudhry, Cengiz Pehlevan
John A. Paulson School of Engineering and Applied Sciences
Center for Brain Science
Kempner Institute for the Study of Natural and Artificial Intelligence
Harvard University
Cambridge, MA 02138
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

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

##  1 Introduction

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

One approach to scaling up systems in a stable and predictable way is to
identify parameterizations of neural networks that give approximately scale-
independent feature updates during training [10, 11, 12]. The mean field
parameterization, commonly referred to as $\mu$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. 1.

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

  2. 2.

We analytically show that the $N\to\infty$ limit requires the $\mu$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. 3.

From the limiting equations, we show that this $N\to\infty$ 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. 4.

To overcome this limitation, we analyze the $\mathcal{H}\to\infty$ 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
$\mathcal{H}\to\infty$ limit leads to concentration of the network’s output
logits and learned residual stream feature kernels, giving deterministic
training dynamics.

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

###  1.1 Related Works

Hron et al. [18] studied the Neural Network Gaussian Process limit of multi-
head self attention in the limit of $\mathcal{H}\to\infty$. 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\to\infty$ using a $\frac{1}{\sqrt{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
$\mathcal{H}\to\infty$ (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^{-\alpha_{L}}$ with
$\alpha_{L}\in[\frac{1}{2},1]$, which has been shown to converge to a limit
not only at initialization [25, 26], but also throughout training in the
feature learning regime [11, 12].

Yang et al. [10] introduced the $\mu$P scaling for attention layers which
multiplies the key/query inner product with $\frac{1}{N}$ rather than the more
commonly used $\frac{1}{\sqrt{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 $\mu$P
transformers learn attention matrices that become approximately consistent
across different heads and model sizes, suggesting that models parameterized
in $\mu$P learn similar representations across scales.

##  2 Parameterizations with Feature Learning Limits

We consider a transformer architecture with $L$ layers, $\mathcal{H}$ heads
per layer, and $N$ dimensional keys/ queries per head. Transformers are often
defined in terms of $d_{\text{model}}=\mathcal{H}d_{\text{head}}=\mathcal{H}N$
which can be increased by scaling the number of heads or the dimension of each
head, where $N$ is often written $d_{\text{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,\mathcal{H},L\to\infty$ limits.

###  2.1 Model Scalings

The network’s output is computed by a depth $L$ recursion through hidden
layers $\ell\in[L]$ starting with the first layer
$\bm{h}^{1}_{\mathfrak{s}}(\bm{x})=\frac{1}{\sqrt{D}}\bm{W}^{0}\bm{x}_{\mathfrak{s}}\in\mathbb{R}^{N\mathcal{H}}$
where $\bm{x}_{\mathfrak{s}}\in\mathbb{R}^{D}$ is the input at spatial/token
position $\mathfrak{s}$. Preactivations in subsequent layers $\bm{h}^{\ell}$
are determined by a forward pass through the residual stream which contains an
attention layer and a MLP layer

$\displaystyle\bm{h}^{\ell+1}_{\mathfrak{s}}=\tilde{\bm{h}}^{\ell}_{\mathfrak{s}}+\frac{\beta_{0}}{L^{\alpha_{L}}}\text{MLP}\left(\tilde{\bm{h}}^{\ell}_{\mathfrak{s}}\right)\
,\
\tilde{\bm{h}}^{\ell}_{\mathfrak{s}}=\bm{h}^{\ell}_{\mathfrak{s}}+\frac{\beta_{0}}{L^{\alpha_{L}}}\text{MHSA}\left(\bm{h}^{\ell}\right)_{\mathfrak{s}}.$
(1)

The constants $\gamma_{0}$ and $\beta_{0}$ control the rate of feature
learning and the effective depth respectively 111The scale of the update to
the residual stream due to each layer will be
$\mathcal{O}(\gamma_{0}\beta_{0}^{2}/L)$. . We will consider
$\alpha_{L}\in[\frac{1}{2},1]$.222If $\alpha_{L}<\frac{1}{2}$ or
$\alpha_{\mathcal{A}}<\frac{1}{2}$ some of the forward pass variables would
diverge at initialization as $L\to\infty$ or $N\to\infty$ respectively. If
$\alpha_{L}>1$ then updates to internal residual blocks will diverge as
$L\to\infty$.  The multi-head self attention layer (MHSA) with pre-layer-
norm333The LN denotes layer-norm which we let be defined in terms of each
vector’s instantaneous mean and variance
$\text{LN}(\bm{h})=\frac{1}{\sqrt{\sigma^{2}+\epsilon}}(\bm{h}-\mu\bm{1})$
where $\mu=\frac{1}{N\mathcal{H}}\bm{1}\cdot\bm{h}$ and
$\sigma^{2}=\frac{1}{N\mathcal{H}}|\bm{h}-\mu\bm{1}|^{2}$. has the form

$\displaystyle\text{MHSA}\left(\bm{h}^{\ell}\right)_{\mathfrak{s}}=\frac{1}{\sqrt{N\mathcal{H}}}\sum_{\mathfrak{h}\in[\mathcal{H}]}\bm{W}^{\ell}_{O\mathfrak{h}}\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}}\
$
$\displaystyle,\quad\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}}=\sum_{\mathfrak{s}^{\prime}\in[\mathcal{S}]}\sigma^{\ell}_{\mathfrak{h}^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\bm{v}^{\ell}_{\mathfrak{h}^{\prime}\mathfrak{s}^{\prime}}$
$\displaystyle\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{V\mathfrak{h}}\bar{\bm{h}}^{\ell}_{\mathfrak{s}}\
$
$\displaystyle,\quad\bar{\bm{h}}^{\ell}_{\mathfrak{s}}=\text{LN}(\bm{h}^{\ell}_{\mathfrak{s}}),$
(2)

where
$\bm{\sigma}^{\ell}_{\mathfrak{h}}\in\mathbb{R}^{\mathcal{S}\times\mathcal{S}}$
are the attention matrices passed through a matrix-valued nonlinearity
$\sigma\left(\bm{\mathcal{A}}^{\ell}_{\mathfrak{h}}\right)$444The nonlinearity
is generally softmax. For autoregressive tasks, it is taken to be causal.. For
a sequence of length $\mathcal{S}$, the pre-attention matrix
$\bm{\mathcal{A}}^{\ell}_{\mathfrak{h}}\in\mathbb{R}^{\mathcal{S}\times\mathcal{S}}$
is defined as

$\displaystyle\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}=\frac{1}{N^{\alpha_{\mathcal{A}}}}\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}\cdot\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime}}^{\ell}\
,\
\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{K\mathfrak{h}}\bar{\bm{h}}^{\ell}_{\mathfrak{s}}\
,\
\bm{q}_{\mathfrak{h}\mathfrak{s}}^{\ell}=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{Q\mathfrak{h}}\bar{\bm{h}}^{\ell}_{\mathfrak{s}}.$
(3)

The exponent $\alpha_{\mathcal{A}}$ will alter the scale of the pre-attention
variables $\mathcal{A}^{\ell}_{\mathfrak{h}}$ at initialization. The input
matrices have shape
$\bm{W}^{\ell}_{V\mathfrak{h}},\bm{W}^{\ell}_{K\mathfrak{h}},\bm{W}^{\ell}_{Q\mathfrak{h}}\in\mathbb{R}^{N\times
N\mathcal{H}}$, while the output matrices have shape
$\bm{W}^{\ell}_{O\mathfrak{h}}\in\mathbb{R}^{N\mathcal{H}\times N}$. All of
the weights
$\bm{W}^{\ell}_{O\mathfrak{h}^{\prime}},\bm{W}^{\ell}_{Q\mathfrak{h}},\bm{W}^{\ell}_{K\mathfrak{h}}$
are initialized with $\Theta(1)$ entries while
$\bm{W}^{\ell}_{K\mathfrak{h}},\bm{W}^{\ell}_{Q\mathfrak{h}}$ have entries of
size $\Theta(N^{1-\alpha_{\mathcal{A}}})$ which ensures that all key and query
$\bm{k},\bm{q}$ vectors are $\Theta(1)$ at initialization. The pre-attention
variables
$\bm{\mathcal{A}}_{\mathfrak{h}}^{\ell}\in\mathbb{R}^{\mathcal{S}\times\mathcal{S}}$
at each head $\mathfrak{h}$ are determined by key
$\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}$ and query
$\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}$ inner products. The MLP
layer consists of two linear layers with an element-wise nonlinearity $\phi$
applied in between, where
$\bm{W}^{\ell,2},\bm{W}^{\ell,1}\in\mathbb{R}^{N\mathcal{H}\times
N\mathcal{H}}$ are initialized with $\Theta(1)$ entries:

$\displaystyle\text{MLP}(\tilde{\bm{h}}^{\ell}_{\mathfrak{s}})=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell,2}\phi\left(\tilde{\bm{h}}^{\ell,1}_{\mathfrak{s}}\right)\
,\
\tilde{\bm{h}}^{\ell,1}_{\mathfrak{s}}=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell,1}\bar{\tilde{\bm{h}}}^{\ell}_{\mathfrak{s}}\
,\
\bar{\tilde{\bm{h}}}^{\ell}_{\mathfrak{s}}=\text{LN}\left(\tilde{\bm{h}}^{\ell}_{\mathfrak{s}}\right).$
(4)

$\mu$P scaling [13, 27, 14, 15] downscales the readout of the last layer
compared to standard and NTK parameterization [28]. Thus, we define the output
of the model as
$f=\frac{1}{\gamma_{0}N\mathcal{H}}\bm{w}^{L}\cdot(\frac{1}{\mathcal{S}}\sum_{\mathfrak{s}}\bm{h}^{L}_{\mathfrak{s}})$555Instead
of pooling over space, outputs $f$ could also carry spatial index
$\mathfrak{s}$ (such as next word prediction). where
$\bm{h}^{L}_{\mathfrak{s}}\in\mathbb{R}^{N\mathcal{H}}$ are the final layer
preactivations at spatial/token position $\mathfrak{s}\in[\mathcal{S}]$.

###  2.2 Learning Rate Scalings

In order to approximately preserve the size of internal feature updates, we
must scale the learning rate $\eta$ appropriately with $(N,\mathcal{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
$\alpha_{L}\in[\frac{1}{2},1]$ and $\alpha_{\mathcal{A}}\in[\frac{1}{2},1]$.
In what follows, we focus on the SGD scaling and theoretically analyze the
$N\to\infty$, $\mathcal{H}\to\infty$, and $L\to\infty$ 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\to\infty$ limit. Example FLAX implementations of these
parameterizations for vision and language modeling transformers are provided
in Appendix B.

Optimizer | Bulk Parameters LR | First/Last Layer Rescale Factor
---|---|---
SGD | $\eta_{0}N\mathcal{H}L^{2\alpha_{L}-1}$ | $L^{\frac{1}{2}-\alpha_{L}}$
Adam | $\eta_{0}N^{-1/2}\mathcal{H}^{-1/2}L^{-1+\alpha_{L}}$ | $L^{1-\alpha_{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 with depth
depending on the parameterization and optimizer.

Our analysis assumes that at each step $t$ of SGD or Adam a mini-batch
$\mathfrak{B}_{t}$ of size $\Theta(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 $\Theta(1)$ sized
dataset.

##  3 Infinite Limits of Learning Dynamics

In this section, we first analyze the $N\to\infty$ limit of training. We find
that for this limit, the $\mu$P rule of $\alpha_{\mathcal{A}}=1$ is necessary
and show that all heads collapse to the same dynamics. To counteract this
effect, we next analyze the $\mathcal{H}\to\infty$ 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
$L\to\infty$ limits. 666Training time, sample size, sequence length/spatial
dimension are all treated as fixed $\Theta(1)$ quantities. .

###  3.1 Dimension-Per-Head $N$

One way of obtaining a well-defined infinite parameter limit of transformers
is to take the $N\to\infty$ 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
$\alpha_{\mathcal{A}}$ are admissible for the pre-attention
$\mathcal{A}=\frac{1}{N^{\alpha_{\mathcal{A}}}}\bm{k}\cdot\bm{q}$. The keys
and queries are uncorrelated at initialization which motivated the original
choice of $\alpha_{\mathcal{A}}=\frac{1}{2}$ [5, 18]. Yang et al. [10] assume
the entries of the key and query vectors move by $\Theta(1)$, implying
$\alpha_{\mathcal{A}}=1$ is necessary since the update to $\bm{k}$ is
correlated to $\bm{q}$ and vice versa. However, it is possible to obtain
$\Theta(1)$ updates to the attention variable for alternative values of
$\alpha_{\mathcal{A}}$ if we choose the change to key (also query) entries
after gradient descent to be $\delta
k_{i}\sim\Theta(N^{-1+\alpha_{\mathcal{A}}})$. We show that this scaling can
approximately preserve optimal hyperparameters across $N$ in Figure 1 (a) and
give similar dynamics under SGD Appendix C. However, as we show in Appendix
D.1.2, any well defined $N\to\infty$ limit of SGD requires
$\alpha_{\mathcal{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\to\infty$ unless initial key and
query weight matrices are downscaled to have variance of order
$\Theta_{N}(1)$.

(a) Hyperparameter Transfer for Various $\alpha_{\mathcal{A}}$ (b) Attention
variance across heads

Figure 1: Increasing dimension-per-head $N$ with heads fixed for
$\alpha_{\mathcal{A}}=\\{1,\frac{1}{2}\\}$. (a) Both $\alpha_{\mathcal{A}}=1$
and $\alpha_{\mathcal{A}}=\frac{1}{2}$ exhibit similar hyperparameter transfer
for vision transformers trained on CIFAR-5M over finite $N$ at
$\mathcal{H}=16$. (b) The variance of attention variables across the different
heads of a vision transformer after training for $2500$ steps on CIFAR-5M. For
$\alpha_{\mathcal{A}}=1$ the variance of attention variables decays at rate
$\mathcal{O}(N^{-2})$ and for $\alpha_{\mathcal{A}}=\frac{1}{2}$ the variance
does not decay with $N$.

In Appendix D, we provide a DMFT analysis of the $N\to\infty$ 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\to\infty$ limit of transformer SGD
training requires taking $\alpha_{\mathcal{A}}=1$ ($\mu$P scaling), even if
key/query updates are allowed to be rescaled to account for their correlation.
The limiting dynamics of training are governed by the residual stream kernel
$H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N\mathcal{H}}\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
and a collection of inner product kernels in each head $\mathfrak{h}$ that
concentrate as $N\to\infty$

$\displaystyle
V^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N}\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\
,\
Q_{\mathfrak{h}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N}\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(5) $\displaystyle
K_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N}\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\
,\
\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\frac{1}{N}\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t),$
(6)

alongside residual-stream gradient kernels and response functions in the sense
of [15, 11]. The NN output logits $f(\bm{x},t)$ evolve deterministically
according to the above kernels as well as kernels for the gradient vectors
$\bm{g}^{\ell}\equiv\gamma_{0}N\mathcal{H}\frac{\partial
f}{\partial\bm{h}^{\ell}}$ which appear in the backward pass. These variables
become identical across heads such that for any
$\mathfrak{h},\mathfrak{h}^{\prime}\in[\mathcal{H}]$,
$\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\mathcal{A}^{\ell}_{\mathfrak{h}^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)$.
All preactivations on the residual stream and key/query/value variables within
a MHSA block are statistically independent across neurons and can be described
by a single scalar stochastic process

$\displaystyle
h^{\ell+1}_{\mathfrak{s}}(\bm{x},t)={h}^{\ell}_{\mathfrak{s}}(\bm{x},t)+\beta_{0}L^{-\alpha_{L}}\tilde{u}^{\ell}_{\mathfrak{s}}(\bm{x},t)+\beta_{0}L^{-\alpha_{L}}u^{\ell+1}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle\qquad+\eta_{0}\gamma_{0}\beta_{0}^{2}L^{-1}\sum_{t^{\prime}<t}\sum_{\mathfrak{s}^{\prime}\in[\mathcal{S}]}\int
d\bm{x}^{\prime}\left[\tilde{C}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\tilde{g}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})+C^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})g^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]$
$\displaystyle
k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=u^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\sum_{t^{\prime}\mathfrak{s}^{\prime}}\int
d\bm{x}^{\prime}C^{k^{\ell}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(7)

where $\tilde{u}^{\ell},u^{\ell},u^{\ell}_{K\mathfrak{h}}$ are Gaussian
processes with covariances $\Phi^{\ell,1},V^{\ell\sigma},H^{\ell}$
respectively. Analogous equations hold for the queries and values. In the
limit, the kernels
$H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<h^{\ell}_{\mathfrak{s}}(\bm{x},t)h^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>$
,
$\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\left<k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right>$,
etc. are computed as averages $\left<\cdot\right>$ over these random
variables. The deterministic kernels $C^{\ell},\tilde{C}^{\ell}$ can also be
expressed in terms of single site averages of residual variables and head
averages of MHSA variables. The kernels
$C^{\ell},\tilde{C}^{\ell},C^{k^{\ell}}$ depend on the precise mini-batches of
data $\mathfrak{B}_{t}$ presented at each step $t$ which we assume are known.

We derive this result using a Martin-Siggia-Rose path integral formalism [29]
for DMFT in Appendix D. Full DMFT equations can be found in Appendix D.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, 30, 17].

##### Collapse of Attention Heads

As $N\to\infty$, multi-head self-attention will effectively compute the same
outputs as a single-head self-attention block. We theoretically derive this
effect in Appendix D.2.1 and demonstrate it empirically in Figure 1 (b). This
experiment shows that in $\alpha_{\mathcal{A}}=1$ ($\mu$P) transformers
trained for $2500$ steps on CIFAR-5M [31], the variance of attention matrices
across heads decreases with $N$. However, we note that if the scaling exponent
is chosen instead as $\alpha_{\mathcal{A}}=\frac{1}{2}$ there is non-
decreasing diversity of attention variables across heads. This result is
consistent with recent empirical findings that attention variables in $\mu$P
transformers converge to the same limiting quantities at large $N$ with
$\mathcal{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.

###  3.2 Infinite Heads

In this section, we take $\mathcal{H}\to\infty$ 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 $\mathcal{H}\to\infty$ limit of SGD training dynamics in a
randomly initialized transformer at any key/query dimension $N$, scaling
exponents $\alpha_{\mathcal{A}},\alpha_{L}$, and any depth $L$ is governed by
head-averaged kernels for pairs of input data $\bm{x},\bm{x}^{\prime}$ at
training times $t,t^{\prime}$ and spatial/token positions
$\mathfrak{s},\mathfrak{s}^{\prime}$ such as

$\displaystyle{V}^{\ell,\sigma}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(8)

which converge to deterministic values as $\mathcal{H}\to\infty$. The
attention variables
$\\{\bm{k}_{\mathfrak{h}}^{\ell}(\bm{x},t),\bm{q}_{\mathfrak{h}}^{\ell}(\bm{x},t),\bm{v}^{\ell}_{\mathfrak{h}}(\bm{x},t),\mathcal{A}_{\mathfrak{h}}^{\ell}(\bm{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 D.3, showing that the joint
distribution of head-averaged dynamical quantities satisfies a large deviation
principle and the limit can be derived as a saddle point of a DMFT action.

To gain intuition for this result, we first examine variables $H^{\ell}$ and
$\mathcal{A}^{\ell}_{\mathfrak{h}}$ at initialization. In Figure 2, we plot
the convergence of a $N=4,L=8$ vision transformer’s residual stream kernel
$H^{\ell}$ to its $\mathcal{H}\to\infty$ limit at rate
$\mathcal{O}(\mathcal{H}^{-1})$ in square error, consistent with perturbative
analysis near the limit [17]. Next, we plot the distribution (over heads) of
$\mathcal{A}_{\mathfrak{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\to\infty$ the distribution of $\mathcal{A}$ will approach a Gaussian with
variance $\Theta(N^{1-2\alpha_{\mathcal{A}}})$.

(a) Initial ${H}^{\ell}$ Kernels, $N=4$ (b) $\mathcal{A}_{\mathfrak{h}}$
Distribution $N=1$ (c) $\mathcal{A}_{\bm{h}}$ Distribution $N=16$

Figure 2: The initial kernels converge as $\mathcal{H}\to\infty$ and are
determined by (possibly non-Gaussian) distributions of
$\mathcal{A}_{\mathfrak{h}}^{\ell}$ over heads in each layer. (a) Convergence
of
${H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime})=\frac{1}{\mathcal{H}N}\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x})\cdot\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime})$
in a $L=8,N=4$ vision transformer at initialization at rate
$\mathcal{O}(\mathcal{H}^{-1})$. (b) The density of
$\mathcal{A}_{\mathfrak{h}}^{\ell}$ entries over heads at fixed spatial
location converges as $\mathcal{H}\to\infty$ but is non-Gaussian for small
$N$. (c) As $N\to\infty$ the initial density of $\mathcal{A}$ approaches a
Gaussian with variance of order $\mathcal{O}(N^{1-2\alpha_{\mathcal{A}}})$.

We then investigate training dynamics as we approach the
$\mathcal{H}\to\infty$ limit. In Figure 3 (a) we show the test loss on
CIFAR-5M as a function of the number of training iterations. The performance
tends improve as $\mathcal{H}$ increases and the model approaches its limit.
In Figure 3 (b) we show that all of the models are converging in function
space by measuring the squared error between finite $\mathcal{H}$ head models
and a proxy for the infinite $\mathcal{H}$ model. Since the
$\mathcal{H}\to\infty$ 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
$\mathcal{H}$ head models converge to the limit proxy at a rate
$\mathcal{O}(\mathcal{H}^{-1})$, but after continued training the convergence
rate weakens. This effect has been observed in $\mu$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 $\mathcal{H}$
kernels [32].

(a) Training Dynamics Varying $\mathcal{H}$ (b) Convergence to
$\mathcal{H}\to\infty$ limit

Figure 3: Approaching the large head limit $\mathcal{H}\to\infty$ in early
portion of SGD dynamics for a vision transformer trained on CIFAR-5M with
$(L,N)=(2,4)$ and
$(\gamma_{0},\beta,\alpha_{\mathcal{A}})=(0.05,4,\frac{1}{2})$ and losses
averaged over $10$ random inits (colored error bars are standard deviations).
(a) As $\mathcal{H}$ increases the loss and the variability over random
initial seed decreases. (b) The mean square difference between output logits
for $\mathcal{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].

###  3.3 Infinite Depth Limits

We next describe the infinite depth limits which depend on the choice of
$\alpha_{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 $\mathcal{H},L\to\infty$ with
$L^{-\alpha_{L}}$ branch scaling with $\alpha_{L}\in[\frac{1}{2},1]$ is
described by a differential equation for residual variables
$h_{\mathfrak{s}}(\tau,t)$ in layer time
$\tau=\lim_{L\to\infty}\frac{\ell}{L}$ for the residual stream

$\displaystyle h_{\mathfrak{s}}(\tau,\bm{x},t)$ $\displaystyle=\beta_{0}\
\delta_{\alpha_{L},\frac{1}{2}}\int_{0}^{\tau}du_{\mathfrak{s}}(\tau^{\prime},\bm{x},t)$
$\displaystyle\qquad+\eta_{0}\gamma_{0}\beta_{0}^{2}\sum_{t^{\prime}<t}\int
d\bm{x}^{\prime}\int_{0}^{\tau}d\tau^{\prime}C_{\mathfrak{s}\mathfrak{s}^{\prime}}(\tau^{\prime},\bm{x},\bm{x}^{\prime},t,t^{\prime})g_{\mathfrak{s}^{\prime}}(\tau^{\prime},\bm{x}^{\prime},t^{\prime})$
(9)

where the Brownian motion term $du_{\mathfrak{s}}(\tau,\bm{x},t)$ survives in
the limit only if $\alpha_{L}=\frac{1}{2}$ and has covariance

$\displaystyle\left<du_{\mathfrak{s}}(\tau,\bm{x},t)du_{\mathfrak{s}^{\prime}}(\tau^{\prime},\bm{x}^{\prime},t^{\prime})\right>=\delta(\tau-\tau^{\prime})d\tau
d\tau^{\prime}\left[\Phi_{\mathfrak{s}\mathfrak{s}^{\prime}}(\tau,\bm{x},\bm{x}^{\prime},t,t^{\prime})+V^{\sigma}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\tau,\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
(10)

and the deterministic kernel
$C_{\mathfrak{s}\mathfrak{s}^{\prime}}(\tau,\bm{x},\bm{x}^{\prime},t,t^{\prime})$
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\to\infty$ limit unless $\alpha_{L}=1$. All response functions are
suppressed at $L\to\infty$ unless $\alpha_{L}=\frac{1}{2}$.

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

  1. 1.

At initialization $t=0$, the only term which contributes to the residual
stream layer dynamics is the integrated Brownian motion
$\int_{0}^{\tau}du(\tau^{\prime})$ which survives at infinite depth for
$\alpha_{L}=\frac{1}{2}$. For $\alpha_{L}=1$ this term disappears in the
limit. The structure of $C(\tau)$ is also modified by additional response
functions at $\alpha_{L}=\frac{1}{2}$ [11] which we show disappear for
$\alpha_{L}=1$.

  2. 2.

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

  3. 3.

A consequence of our large $N$ and large $L$ result is that the $N,L\to\infty$
limit with $\alpha_{L}=\frac{1}{2}$ (the parameterization studied by [11, 12])
would lead to
$\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{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^{\ell}$ updating from feature learning.

We can see some numerical evidence for the first of these effects in Figure 4
(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\to\infty$ for $\alpha_{L}=1$ since
the preactivation vectors lose variance as we discuss in Appendix D.4,
resulting in slower initial training. However, we note that the final learned
feature kernels are quite similar after enough training.

(a) Training Dynamics (b) Initial and Final Residual Stream Pooled Kernels (c)
Spatial Kernels for Single Sample (d) Attention Distributions Before and After
Training

Figure 4: Initial and final representations are converging as model scale
increases after one pass of training on the full CIFAR-5M with SGD+momentum.
The base model is a $(N,\mathcal{H},L)=(16,16,4)$ and
$(\alpha_{\mathcal{A}},\alpha_{L},\beta_{0},\gamma_{0})=(1,1,4,0.1)$. (a) The
test loss dynamics for one pass through CIFAR-5M. The dynamics are very
similar across different $\mathcal{H}$ but the early dynamics are changed for
large $L$, consistent with our theory. (b) The initial and final feature
kernels after spatial pooling at the last layer of the residual stream. The
initial kernel at large $L$ is quite different for $\alpha_{\mathcal{A}}=1$
due to suppression of Brownian motion on the forward pass, which we explain in
Section 3.3. (c) The residual stream kernel across pairs of spatial positions
for a single randomly chosen input sample. (d) The distribution of attention
entries across heads at a fixed pair of spatial locations and data point. The
initial variance of $\mathcal{A}$ decreases for $\alpha_{\mathcal{A}}=1$ but
the update is roughly consistent across $N$. For
$\alpha_{\mathcal{A}}=\frac{1}{2}$ both initial and final distributions for
$\mathcal{A}_{\mathfrak{h}}$ are consistent across $N$.

In summary, our results indicate that the $\alpha_{L}=1$ parameterization is
the one that allows attention layers to actually be learned in the limit
$L\to\infty$, but that this parameterization leads to a less structured kernel
at initialization.

##  4 Experiments in Realistic Settings

(a) Training dynamics of LLM on C4 (b) Kernels: final token across samples (c)
Kernels: tokens within single sample

Figure 5: Training dynamics and initial/final representations of decoder only
language models trained on C4 converge with increasing model scale. The base
model has $(N,\mathcal{H},L)=(8,8,4)$ and
$(\alpha_{L},\beta_{0},\gamma_{0})=(1,4,0.25)$ and
$\alpha_{\mathcal{A}}\in\\{1,\frac{1}{2}\\}$. (a) Train loss dynamics after
$10000$ steps on C4 using Adam optimizer. The dynamics improve consistently
when scaling $\mathcal{H}$ for both values of $\alpha_{\mathcal{A}}$, with
slight benefit to $\alpha_{\mathcal{A}}=\frac{1}{2}$. Scaling $N$ reveals a
significant advantage to setting $\alpha_{\mathcal{A}}=\frac{1}{2}$. Scaling
$L$ provides little improvement for either parameterization of
$\alpha_{\mathcal{A}}$. (b) Initial and final residual stream kernels for the
final token across samples for $\text{Base},\mathcal{H}=128$, $N=128$, and
$L=64$ models. The first row is at initialization. The second and third rows
are after training with $\alpha_{\mathcal{A}}\in\\{1,\frac{1}{2}\\}$
respectively. (c) Initial and final feature kernels across pairs of tokens for
a single randomly chosen input sample. Note both types of kernels are
identical across $\alpha_{\mathcal{A}}$ except for a slight difference at
large $N$.

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 [33] 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 to investigate realistic training
dynamics of our LLM as we take the $N,L,\mathcal{H}\to\infty$ limits. Training
details are provided in Appendix E

In Figure 5 (a), we sweep over each of the model dimensions independently for
each parameterization of $\alpha_{\mathcal{A}}\in\\{1,\frac{1}{2}\\}$ on the
left and right respectively. For fixed $N$ and $L$, scaling $\mathcal{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
$\alpha_{\mathcal{A}}=\frac{1}{2}$. For fixed $\mathcal{H}$ and $L$, scaling
$N$ provides a similar increase in performance to scaling heads in when
$\alpha_{\mathcal{A}}=1$, but a substantial increase when
$\alpha_{\mathcal{A}}=\frac{1}{2}$. This is in line with our predictions in
Section 3.1 about the benefits of diversity across attention heads. Next, for
fixed $N$ and $\mathcal{H}$, scaling $L$ provides little to no benefit in
either parameterization as predicted in Section 3.3. 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
$\alpha_{\mathcal{A}}$, except for a slight difference for large $N$.
Furthermore, they are extremely similar for large $N$ and large $\mathcal{H}$.

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

##  5 Discussion

This paper provided analysis of the infinite head, depth and key/query
dimension limits of transformer training in the feature learning regime. We
showed that feature learning in $\mu$P multi-head transformers in the limit of
$N\to\infty$ collapses to single-head self-attention. At finite $N$ and
infinite heads $\mathcal{H}\to\infty$ 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 $\tau$. The depth scaling
that maintains feature learning within all MHSA blocks ($\alpha_{L}=1$) causes
the initial kernel to lose structure from the initialization as $L\to\infty$,
but allows learning of the self-attention variables, whereas the depth scaling
that preserves structure from initialization ($\alpha_{L}=\frac{1}{2}$) leads
to static layers.

##### Limitations and Future Directions

Currently exact theoretical analysis of the limit is focused on SGD (and can
be easily extended to SGD+momentum [15]) while Adam is currently only reasoned
with rough scaling arguments rather than an exact theoretical description of
the limit. Since Adam is most commonly used to train transformers, a theory of
the limiting dynamics of Adam in Transformers would be an important future
extension. In addition, while we provide an exact asymptotic description of
network training, the limiting equations are compute intensive for realistic
settings which is why we focus our empirical investigations on training large
width networks in the appropriate parameterizations. Lastly our techniques
assume that the number of training steps is fixed as the scaling parameters of
interest $(N,\mathcal{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, 32, 34]. In this regime,
harmful finite model-size effects become significant and comparable to the
finite training horizon [10, 16, 17, 32]. 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.

## Acknowledgements and Disclosure of Funding

BB would like to thank Alex Atanasov, Jacob Zavatone-Veth, Lorenzo Noci, Mufan
Bill Li, Boris Hanin, Alex Damian, Eshaan Nichani for inspiring conversations.
We would also like to thank Alex Atanasov and Jacob Zavatone-Veth for useful
comments on an earlier version of this manuscript. BB is supported by a Google
PhD fellowship. HC was supported by the GFSD Fellowship, Harvard GSAS Prize
Fellowship, and Harvard James Mills Peirce Fellowship. CP was supported by NSF
Award DMS2134157 and NSF CAREER Award IIS2239780. CP is further supported by a
Sloan Research Fellowship. This work has been made possible in part by a gift
from the Chan Zuckerberg Initiative Foundation to establish the Kempner
Institute for the Study of Natural and Artificial Intelligence. The
computations in this paper were run on the FASRC Cannon cluster supported by
the FAS Division of Science Research Computing Group at Harvard University.

## References

  * Liu et al. [2021] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 10012–10022, 2021.
  * Han et al. [2022] Kai Han, Yunhe Wang, Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, et al.  A survey on vision transformer.  _IEEE transactions on pattern analysis and machine intelligence_ , 45(1):87–110, 2022.
  * Dosovitskiy et al. [2020] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Dehghani et al. [2022] Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab, Matthias Minderer, and Yi Tay.  Scenic: A jax library for computer vision research and beyond.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 21393–21398, 2022.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Kaplan et al. [2020] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Hoffmann et al. [2022] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Achiam et al. [2023] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Yang et al. [2021] Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao.  Tuning large neural networks via zero-shot hyperparameter transfer.  _Advances in Neural Information Processing Systems_ , 34:17084–17097, 2021.
  * Bordelon et al. [2024a] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan.  Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit.  In _The Twelfth International Conference on Learning Representations_ , 2024a.  URL https://openreview.net/forum?id=KZJehvRKGD.
  * Yang et al. [2023] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou.  Feature learning in infinite depth neural networks.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Mei et al. [2019] Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.  In _Conference on Learning Theory_ , pages 2388–2464. PMLR, 2019.
  * Yang and Hu [2021] Greg Yang and Edward J Hu.  Tensor programs iv: Feature learning in infinite-width neural networks.  In _International Conference on Machine Learning_ , pages 11727–11737. PMLR, 2021.
  * Bordelon and Pehlevan [2022a] Blake Bordelon and Cengiz Pehlevan.  Self-consistent dynamical field theory of kernel evolution in wide neural networks.  _Advances in Neural Information Processing Systems_ , 35:32240–32256, 2022a.
  * Vyas et al. [2023] Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan.  Feature-learning networks are consistent across widths at realistic scales, 2023.
  * Bordelon and Pehlevan [2023] Blake Bordelon and Cengiz Pehlevan.  Dynamics of finite width kernel and prediction fluctuations in mean field neural networks.  _arXiv preprint arXiv:2304.03408_ , 2023.
  * Hron et al. [2020] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak.  Infinite attention: Nngp and ntk for deep attention networks.  In _International Conference on Machine Learning_ , pages 4376–4386. PMLR, 2020.
  * Dinan et al. [2023] Emily Dinan, Sho Yaida, and Susan Zhang.  Effective theory of transformers at initialization, 2023.
  * Dong et al. [2021] Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: Pure attention loses rank doubly exponentially with depth.  In _International Conference on Machine Learning_ , pages 2793–2803. PMLR, 2021.
  * Noci et al. [2022] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi.  Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.  _Advances in Neural Information Processing Systems_ , 35:27198–27211, 2022.
  * He and Hofmann [2023] Bobby He and Thomas Hofmann.  Simplifying transformer blocks.  _arXiv preprint arXiv:2311.01906_ , 2023.
  * Cowsik et al. [2024] Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli.  Geometric dynamics of signal propagation predict trainability of transformers, 2024.
  * Noci et al. [2024] Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy.  The shaped transformer: Attention models in the infinite depth-and-width limit.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Hayou [2023] Soufiane Hayou.  On the infinite-depth limit of finite-width neural networks.  _Transactions on Machine Learning Research_ , 2023.  ISSN 2835-8856.  URL https://openreview.net/forum?id=RbLsYz1Az9.
  * Cirone et al. [2023] Nicola Muca Cirone, Maud Lemercier, and Cristopher Salvi.  Neural signature kernels as infinite-width-depth-limits of controlled resnets.  _arXiv preprint arXiv:2303.17671_ , 2023.
  * Chizat and Bach [2018] Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  _Advances in neural information processing systems_ , 31, 2018.
  * Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Martin et al. [1973] Paul Cecil Martin, ED Siggia, and HA Rose.  Statistical dynamics of classical systems.  _Physical Review A_ , 8(1):423, 1973.
  * Bordelon and Pehlevan [2022b] Blake Bordelon and Cengiz Pehlevan.  The influence of learning rule on representation dynamics in wide neural networks.  _arXiv preprint arXiv:2210.02157_ , 2022b.
  * Nakkiran et al. [2020] Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi.  The deep bootstrap framework: Good online learners are good offline generalizers.  _arXiv preprint arXiv:2010.08127_ , 2020.
  * Bordelon et al. [2024b] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan.  A dynamical model of neural scaling laws, 2024b.
  * Raffel et al. [2020] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of machine learning research_ , 21(140):1–67, 2020.
  * Alabdulmohsin et al. [2024] Ibrahim M Alabdulmohsin, Xiaohua Zhai, Alexander Kolesnikov, and Lucas Beyer.  Getting vit in shape: Scaling laws for compute-optimal model design.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Littwin and Yang [2022] Etai Littwin and Greg Yang.  Adaptive optimization in the $\infty$-width limit.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Li et al. [2022] Mufan Li, Mihai Nica, and Dan Roy.  The neural covariance sde: Shaped infinite depth-and-width networks at initialization.  _Advances in Neural Information Processing Systems_ , 35:10795–10808, 2022.

## Appendix

##  Appendix A Additional Figures

(a) Vary $\mathcal{H}$ (b) Vary $N$ (c) Vary $L$

Figure 6: One pass training on CIFAR-5M with vision transformers with the
setting of Figure 4.

Figure 7: Examples of initial and learned kernels in final residual stream
layer with various extrapolations of a base vision transformer model with
$(\mathcal{H},N,L)=(16,16,4)$ trained on CIFAR-5M.

Figure 8: Spatial kernels for a single test point before and after training
across $\mathcal{H},N,L$ values.

(a) $N$ Scaling with $\alpha=1$ (b) $N$ Scaling with $\alpha=\frac{1}{2}$

Figure 9: Early training dynamics on CIFAR-5M in vision transformer with
different dimension-per-head $N$ with heads fixed at $\mathcal{H}=4$ for
$\alpha_{\mathcal{A}}=\\{1,\frac{1}{2}\\}$.

##  Appendix B Implementations for Vision and Causal Language Modeling
Transformers

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

We start by defining a fixed layernorm operation

⬇

1from flax import linen as nn

2import jax.numpy as jnp

3

4

5class LN_Fixed(nn.Module):

6

7 eps: jnp.float32 = 1.0e-6

8 @nn.compact

9

10 def __call__(self, x):

11

12 features = x.shape[-1] # number of features

13 mean = jnp.mean( x , axis = -1 ) # mean of x

14 var = jnp.var( x , axis = -1 ) # var of x

15 out = (x - mean[:,:,jnp.newaxis] ) / jnp.sqrt( var[:,:,jnp.newaxis] +
self.eps )

16 return out

The MHSA layer is implemented as the following where scale_exp represents
$\alpha_{\mathcal{A}}$.

⬇

1# MHSA attention layer

2from einops import rearrange

3

4class Attention(nn.Module):

5 """Multi-head Self-Attention Layer"""

6 scale_exp: jnp.float32

7 dim: int

8 heads: int

9

10 def setup(self):

11

12 self.c = 1.5 - self.scale_exp # exponent for the scale factor

13 kif_qk = nn.initializers.normal(stddev = self.dim**(self.c - 0.5) ) #
possible scaling with N

14 kif_v = nn.initializers.normal(stddev = 1.0 ) # O_N(1) entries

15 # computes key, query, value

16 self.qk_layer = nn.Dense(features = 2 * self.heads * self.dim, kernel_init
= kif_qk, use_bias = False)

17 self.v_layer = nn.Dense(features = self.heads * self.dim, kernel_init =
kif_v, use_bias = False)

18 self.out_layer = nn.Dense(features = self.heads * self.dim, kernel_init =
kif_v, use_bias = False)

19 return

20

21 def __call__(self,inputs):

22

23 qk = self.qk_layer(inputs) / self.heads**(0.5) / self.dim**(self.c)

24 qk = rearrange( qk, ’b l (h d) -> b h l d’ , h = self.heads) # (batch,
heads, loc, d )

25 q,k = jnp.split(qk, 2, axis = -1) # gives q, k each of shape ( batch,
heads, loc, d )

26

27 v = self.v_layer(inputs) / jnp.sqrt( inputs.shape[-1] )

28 v = rearrange(v, ’b l (h d) -> b h l d’, h = self.heads)

29 A = self.dim**(-self.scale_exp) * jnp.einsum(’ijkl,ijml->ijkm’, q, k) #
batch x heads x loc x loc

30 sigma_A = softmax( A, axis=-1 )

31 out = jnp.einsum(’ijkl,ijlm->ijkm’, sigma_A, v) # (batch, head, loc, d)

32 out = rearrange(out, ’b h l d -> b l (h d)’)

33 out = self.out_layer(out) / jnp.sqrt( out.shape[-1] )

34 return out

The two layer MLP block is implemented as the following with
$\phi=\text{gelu}$ nonlinearity.

⬇

1class MLP_Block(nn.Module):

2 """Two Layer MLP Block"""

3 features: int

4

5 @nn.compact

6 def __call__(self,x):

7 N = self.features

8 kif = nn.initializers.normal(stddev = 1.0) # O_N(1) entries

9 h = nn.Dense(features = N, kernel_init = kif, use_bias = False)(x) /
jnp.sqrt(N)

10 h = nn.gelu(h)

11 h = nn.Dense(features = N, kernel_init = kif, use_bias = False)(h) /
jnp.sqrt(N)

12 return h

We also allow for a trainable positional encoding matrix.

⬇

1

2class PositionalEncoding(nn.Module):

3 """Trainable Positional Encoding"""

4 d_model : int # Hidden dimensionality of the input.

5 max_len : int # Maximum length of a sequence to expect.

6 scale: jnp.float32 # scale parameter for initialization

7

8 def setup(self):

9 # Create matrix of [SeqLen, HiddenDim] representing the positional encoding
for max_len inputs

10 self.pos_embedding = self.param(’pos_embedding’,

11 nn.initializers.normal(stddev = self.scale),

12 (1, 1+self.max_len, self.d_model))

13

14 def __call__(self, x, train=True):

15 B,T,_ = x.shape

16 x = x + self.pos_embedding[:,:T] / self.scale

17 return x

Each residual block is implemented as the following. Below we show the
$\alpha_{L}=1$ implementation.

⬇

1# Residual Block

2class ResidBlock(nn.Module):

3

4 dim: int

5 heads: int

6 features: int

7 L: int

8 scale_exp: jnp.float32 = 1.0

9 beta: jnp.float32 = 4.0

10

11 @nn.compact

12 def __call__(self,x):

13 h = LN_Fixed()(x)

14 h = Attention(dim = self.dim, scale_exp = self.scale_exp, heads =
self.heads)( h )

15 x = x + self.beta / self.L * h

16 h = LN_Fixed()(x)

17 h = MLP_Block(features = self.features)(h)

18 x = x + self.beta / self.L * h

19 return x

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

⬇

1

2class VIT(nn.Module):

3

4 "simple VIT model with "

5 dim: int

6 heads: int

7 depth: int

8 patch_size: int

9 scale_exp: jnp.float32 = 1.0

10 adam_scale: int = 0.0

11 beta: jnp.float32 = 4.0

12

13 @nn.compact

14 def __call__(self, x):

15 d_model = self.heads * self.dim

16 L = self.depth

17 D = 3

18

19 # patchify images

20 x = rearrange(x, ’b (w p1) (h p2) c -> b (w h) (p1 p2 c)’, p1 =
self.patch_size, p2 = self.patch_size) # (batch, loc, patch_ch_dim )

21

22 kif_first= nn.initializers.normal(stddev = d_model**(-0.5*self.adam_scale)
* (L/self.beta)**(0.5 * (1.0-self.adam_scale)) ) # O_N(1) entries

23 kif = nn.initializers.normal( stddev = 1.0 ) # O_N(1) entries

24 kif_last = nn.initializers.normal(stddev = (L/self.beta)**(0.5 *
(1-self.adam_scale) ) )

25

26 # read-in weights

27 x = (L/self.beta)**(-0.5 * (1.0-self.adam_scale))*d_model**(0.5 *
self.adam_scale) * nn.Dense(features = N, kernel_init = kif_first, use_bias =
False)(x) / jnp.sqrt( D * self.patch_size**2 )

28

29 # positional encoding

30 x = PositionalEncoding(d_model = d_model, max_len =
(32//self.patch_size)**2, scale =
d_model**(-0.5*self.adam_scale)*(L/self.beta)**(0.5 *
(1.0-self.adam_scale)))(x)

31

32 # residual stream with pre-LN

33 for l in range(self.depth):

34 x = ResidBlock(dim = self.dim, heads = self.heads,
scale_exp=self.scale_exp, features = d_model, beta=self.beta, L = L)(x)

35

36 # last norm layer

37 x = LN_Fixed()(x)

38 # pool over spatial dimension

39 x = x.mean(axis = 1) # (batch, d_model)

40 x = (L/self.beta)**(-0.5*(1-self.adam_scale)) * nn.Dense(features = 10,
use_bias = False, kernel_init = kif_last)(x) /
d_model**(1.0-0.5*self.adam_scale) # for mean field scaling

41 return x

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

⬇

1

2class Causal_Attention(nn.Module):

3

4 scale_exp: jnp.float32

5 dim: int

6 heads: int

7 qk_ln: bool = True

8

9 def setup(self):

10

11 self.c = 1.5 - self.scale_exp # exponent for the scale factor

12 kif_qk = nn.initializers.normal(stddev = self.dim**(self.c - 0.5) ) #
possibly needs to be scaled with N

13 kif_v = nn.initializers.normal(stddev = 1.0 ) # O_N(1) entries

14 # computes key, query, value

15 self.qk_layer = nn.Dense(features = 2 * self.heads * self.dim, kernel_init
= kif_qk, use_bias = False)

16 self.v_layer = nn.Dense(features = self.heads * self.dim, kernel_init =
kif_v, use_bias = False)

17 self.out_layer = nn.Dense(features = self.heads * self.dim, kernel_init =
kif_v, use_bias = False)

18 return

19

20 def __call__(self,inputs):

21

22 qk = self.qk_layer(inputs) / self.heads**(0.5) / self.dim**(self.c) #
(batch, loc, 3*h*d)

23 qk = rearrange( qk, ’b l (h d) -> b h l d’ , h = self.heads) # (batch,
heads, loc, d )

24 q,k = jnp.split(qk, 2, axis = -1) # gives q, k each of shape ( batch,
heads, loc, d )

25

26 v = self.v_layer(inputs) / jnp.sqrt( inputs.shape[-1] )

27 v = rearrange(v, ’b l (h d) -> b h l d’, h = self.heads)

28

29 A = 1.0/ self.dim**(self.scale_exp) * jnp.einsum(’ijkl,ijml->ijkm’, q, k) #
batch x heads x loc x loc

30 exp_A = jnp.einsum(’ijkl,kl->ijkl’, jnp.exp(A),
jnp.tril(jnp.ones((v.shape[2], v.shape[2]))))

31 phi_A = exp_A / exp_A.sum(axis = -1)[:,:,:,jnp.newaxis]

32

33 out = jnp.einsum(’ijkl,ijlm->ijkm’, phi_A, v) # (batch, head, loc, d)

34 out = rearrange(out, ’b h l d -> b l (h d)’)

35 out = self.out_layer(out) / jnp.sqrt( out.shape[-1] )

36 return out

37

38

39class LM_Transformer(nn.Module):

40 """A simple Decoder only transformer"""

41

42 dim: int

43 heads: int

44 depth: int

45 scale_exp: jnp.float32

46 adam_scale: int

47 beta: jnp.float32

48 VOCAB_SIZE: int

49

50 @nn.compact

51 def __call__(self, x, train = True):

52 d_model = self.heads * self.dim

53 L = self.depth

54 kif_first = nn.initializers.normal(stddev = d_model**(-0.5*self.adam_scale)
* (L/self.beta)**(0.5 * (1-self.adam_scale) ) ) # O(1) entries

55 kif0 = nn.initializers.normal(stddev = 0.0 )

56 kif = nn.initializers.normal(stddev = 1.0) # O(1) entries

57 kif_last = nn.initializers.normal(stddev = (L/self.beta)**(0.5 *
(1-self.adam_scale)) * d_model**(-0.5*self.adam_scale) )

58

59 # embed the batch x sequence integers to

60 x = (L/self.beta)**( -0.5 * (1-self.adam_scale) )* d_model**(0.5 *
self.adam_scale) * nn.Embed(self.VOCAB_SIZE, d_model, embedding_init =
kif_first)(x) # batch x seq len x N

61

62 x = PositionalEncoding(d_model = d_model, scale =
d_model**(-0.5*self.adam_scale) * (L/self.beta)**(0.5 *(1-self.adam_scale))
)(x)

63

64 for l in range(self.depth):

65 h = LN_Fixed()(x)

66 x = x + self.beta/L * Causal_Attention(dim = self.dim, scale_exp =
self.scale_exp, heads = self.heads)(h)

67 h = LN_Fixed()(x)

68 x = x + self.beta/L * MLP_Block(features = d_model)(h)

69

70 x = LN_Fixed()(x)

71 x = (L/self.beta)**(-0.5 * (1 - self.adam_scale ) ) * nn.Dense(features =
self.VOCAB_SIZE, use_bias = True, kernel_init = kif0)(x) /
d_model**(1.0-0.5*self.adam_scale) # for mean field scaling

72 return x

##  Appendix C Simple Heuristic Scaling Analysis

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

###  C.1 Learning Rate Scalings

We show that the correct learning rate scaling for SGD is
$\eta=\eta_{0}N\mathcal{H}L^{2\alpha_{L}-1}$. For Adam, the learning rate
should be scaled as
$\eta=\eta_{0}N^{-1/2}\mathcal{H}^{-1/2}L^{-1+\alpha_{L}}$.

Optimizer | Bulk Parameters LR | First Layer Rescale Factor
---|---|---
SGD | $\eta_{0}N\mathcal{H}L^{2\alpha_{L}-1}$ | $L^{-\frac{1}{2}-\alpha}$
Adam | $\eta_{0}N^{-1/2}\mathcal{H}^{-1/2}L^{-1+\alpha_{L}}$ | $L^{1-\alpha_{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.

###  C.2 Heuristic Analysis of Feature Changes Under SGD

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

$\displaystyle\delta\bm{W}^{\ell}_{O\mathfrak{h}}\sim\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{g}^{\ell+1}\bm{v}^{\ell\top}_{\mathfrak{h}}$
(11)

where $\bm{g}^{\ell+1}\in\mathbb{R}^{N\mathcal{H}}$ and
$\bm{v}^{\ell}_{\mathfrak{h}}\in\mathbb{R}^{N}$ have $\Theta(1)$ entries.
Thus, computing a perturbation to the forward pass we find

$\displaystyle\delta\bm{h}^{\ell+1}=\delta\bm{h}^{\ell}+\frac{1}{L^{\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\left(\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{g}^{\ell+1}\bm{v}^{\ell\top}_{\mathfrak{h}}\right)\bm{v}^{\ell}_{\mathfrak{h}}$
$\displaystyle=$
$\displaystyle\delta\bm{h}^{\ell}+\frac{1}{L}\left[\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}}\bm{v}^{\ell}_{\mathfrak{h}}\cdot\bm{v}^{\ell}_{\mathfrak{h}}\right]\bm{g}^{\ell+1}$
(12)

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

For the Attention variables, we note that the

$\displaystyle\delta\bm{W}^{\ell}_{K\mathfrak{h}}\sim\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{q}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}\
,\
\delta\bm{W}^{\ell}_{Q\mathfrak{h}}\sim\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{k}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}$
(13)

where $\bm{q}_{\mathfrak{h}},\bm{k}^{\ell}_{\mathfrak{h}}\in\mathbb{R}^{N}$
are the query and key for head $\mathfrak{h}$ and
$\bm{h}\in\mathbb{R}^{N\mathcal{H}}$ is the residual stream preactivation. We
can thus compute the changes to the keys and queries due to changes in their
associated weights

$\displaystyle\delta\bm{k}_{\mathfrak{h}}^{\ell}=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\left(\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{q}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}\right)\bm{h}^{\ell}=\frac{1}{L^{1-\alpha_{L}}N^{1-\alpha_{\mathcal{A}}}}\bm{q}^{\ell}_{\mathfrak{h}}H^{\ell}$
$\displaystyle\delta\bm{q}_{\mathfrak{h}}^{\ell}=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\left(\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{k}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}\right)\bm{h}^{\ell}=\frac{1}{L^{1-\alpha_{L}}N^{1-\alpha_{\mathcal{A}}}}\bm{q}^{\ell}_{\mathfrak{h}}H^{\ell}.$
(14)

where
$H^{\ell}=\frac{1}{N\mathcal{H}}\bm{h}^{\ell}\cdot\bm{h}^{\ell}\sim\Theta(1)$.
Combining these changes , we find the following update to the pre-Attention
variables
$\mathcal{A}^{\ell}_{\mathfrak{h}}=\frac{1}{N^{\alpha_{\mathcal{A}}}}\bm{k}_{\mathfrak{h}}^{\ell}\cdot\bm{q}_{\mathfrak{h}}^{\ell}$

$\displaystyle\delta\mathcal{A}^{\ell}_{\mathfrak{h}}$
$\displaystyle=\frac{1}{L^{1-\alpha_{L}}N}\bm{q}^{\ell}_{\mathfrak{h}}\cdot\bm{q}^{\ell}_{\mathfrak{h}}H^{\ell}+\frac{1}{L^{1-\alpha_{L}}N}\bm{k}^{\ell}_{\mathfrak{h}}\cdot\bm{k}^{\ell}_{\mathfrak{h}}H^{\ell}+\frac{1}{L^{2-2\alpha_{L}}N^{2-2\alpha_{\mathcal{A}}}}\mathcal{A}^{\ell}_{\mathfrak{h}}(H^{\ell})^{2}$
$\displaystyle=\Theta(L^{-1+\alpha_{L}}),$ (15)

since
$\frac{1}{N}\bm{k}\cdot\bm{k},\frac{1}{N}\bm{q}\cdot\bm{q}\sim\Theta(1)$. This
update to the attention variable due to changes in
$\bm{W}^{\ell}_{K},\bm{W}^{\ell}_{Q}$ will clearly die out as $L\to\infty$
unless $\alpha_{L}=1$.

###  C.3 Heuristic Analysis of Feature Changes Under Adam

For Adam, the gradient of each individual parameter entry is approximately
normalized by its scale [35]. Thus the learning rate $\eta$ sets the size of
the updates. This is why we scale the learning rate as
$\eta=\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}$ which gives the same
scale updates to the weights as SGD

$\displaystyle\delta\bm{W}^{\ell}_{O\mathfrak{h}}\approx\eta\
\bm{g}^{\ell+1}\bm{v}^{\ell\top}_{\mathfrak{h}}=\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{g}^{\ell+1}\bm{v}^{\ell\top}_{\mathfrak{h}}$
(16)

Again computing the correction to the forward pass we find

$\displaystyle\delta\bm{h}^{\ell+1}=\delta\bm{h}^{\ell}+\frac{1}{L^{\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\left(\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{g}^{\ell+1}\bm{v}^{\ell\top}_{\mathfrak{h}}\right)\bm{v}^{\ell}_{\mathfrak{h}}$
$\displaystyle=$
$\displaystyle\delta\bm{h}^{\ell}+\frac{1}{L}\left[\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}}\bm{v}^{\ell}_{\mathfrak{h}}\cdot\bm{v}^{\ell}_{\mathfrak{h}}\right]\bm{g}^{\ell+1}=\Theta(1)$
(17)

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

$\displaystyle\delta\bm{W}^{\ell}_{K\mathfrak{h}}\sim\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{q}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}\
,\
\delta\bm{W}^{\ell}_{Q\mathfrak{h}}\sim\frac{1}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\bm{k}^{\ell}_{\mathfrak{h}}\bm{h}^{\ell\top}$
(18)

We can therefore follow the identical argument to identify the scale of the
change to the pre-attention variables
$\delta\mathcal{A}^{\ell}_{\mathfrak{h}}=\Theta(L^{1-\alpha_{L}})$.

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

Any parameterization with $\alpha_{N}\in[\frac{1}{2},1]$ will cause all
updates to $\mathcal{A}_{\mathfrak{h}}^{\ell}$ and entries of
$\bm{h}^{\ell+1}$ to be $\Theta_{N,H,L}(1)$ across finite $N$. The entries of
$\bm{q}$ and $\bm{k}$ only move by $\Theta_{N}(1)$ if $\alpha_{\mathcal{A}}=1$
($\mu$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$, $\mathcal{A}$, $f$, etc) rather than
preactivation or key/query vector entries themselves. Table 3 summarizes two
example values of $\alpha_{\mathcal{A}}$ which are of special interest for
their $N\to\infty$ limits.

| Variance of $\mathcal{A}(0)$ | Update to $\mathcal{A}$ | Update to $\bm{k},\bm{q}$ Entries
---|---|---|---
$\alpha_{\mathcal{A}}=1$ ($\mu$P) | $\Theta(N^{-1})$ | $\Theta(1)$ | $\Theta(1)$
$\alpha_{\mathcal{A}}=\frac{1}{2}$ | $\Theta(1)$ | $\Theta(1)$ | $\Theta(N^{-\frac{1}{2}})$
Table 3: Two interesting choices of scaling for the attention layer exponent
$\alpha_{\mathcal{A}}$ which give approximately constant updates to the
attention matrices $\mathcal{A}_{\mathfrak{h}}$. The $\mu$P scaling
$\alpha_{\mathcal{A}}=1$ causes the entries of the key/query vector entries to
move non-negligibly but causes all heads to be identical (and all
$\mathcal{A}=0$) at initialization. Scaling instead with
$\alpha_{\mathcal{A}}=\frac{1}{2}$ causes the $\mathcal{A}$ variables to be
random but still non-negligibly updated under training.

The choice $\alpha_{\mathcal{A}}=\frac{1}{2}$ allows the variance of
$\mathcal{A}_{\mathfrak{h}}^{\ell}$ to be constant size as a function of $N$
while also enabling learning of these variables. We verify these scalings in
Figure 10.

(a) Change in $\bm{k}$ Entries (SGD) (b) Change in $\mathcal{A}$ Entries (SGD)
(c) Change in $\bm{k}$ Entries (Adam) (d) Change in $\mathcal{A}$ Entries
(Adam)

Figure 10: The update to (a) key $\bm{k}_{\mathfrak{h}}$ entries and (b) pre-
attention variables $\mathcal{A}_{\mathfrak{h}}$ after $t$ steps of gradient
descent for scaling exponents $\alpha_{\mathcal{A}}\in\\{1,\frac{1}{2}\\}$. At
the first step of SGD, the updates to the keys and attention variables are
suppressed due to a lack of correlation between $\bm{W}_{O}$ and the gradient
$\frac{\partial f}{\partial\tilde{\bm{h}}}$. After training for multiple
steps, this correlation increases and non-negligible updates to the attention
variables occur. (c)-(d) The same but for the Adam optimizer with our proposed
parameterization.

##  Appendix D DMFT Analysis for Transformers

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

###  D.1 Deriving the DMFT Action

In this section we will derive the limiting equations of motion for stochastic
gradient descent in the $\mathcal{H}\to\infty$ limit. We start by defining the
loss function which we aim to minimize

$\displaystyle\mathcal{L}=\int d\bm{x}\ p(\bm{x})\ \ell[f(\bm{x})]$ (19)

where $p(\bm{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 $\Delta(x,t)\equiv-\frac{\partial\ell[f(\bm{x},t)]}{\partial
f(\bm{x},t)}$ represent the error signal on datapoint $\bm{x}$. At each step
of training $t$ a batch of examples
$\mathfrak{B}_{t}=\\{\bm{x}_{1}(t),...,\bm{x}_{|\mathfrak{B}_{t}|}(t)\\}$ is
generated and used to estimate a gradient for SGD. In what follows, we let
$\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t}}$ represent averages over the
minibatch at time $t$. We emphasize here that the batches $\mathfrak{B}_{t}$
are assumed to be given or fixed and are not averaged over as random draws
from $p(\bm{x})$, but rather our expectation simply denotes the empirical mean
over the minibatch at time $t$

$\displaystyle\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t}}[f(\bm{x})]=\frac{1}{|\mathfrak{B}_{t}|}\sum_{\bm{x}\in\mathfrak{B}_{t}}f(\bm{x}).$
(20)

We start by expressing again the forward pass equations for each layer. To
make this analysis more compressed while still capturing all of the
interesting aspects, we will first compute the equations of motion in the
absence of MLP layers (which were analyzed in prior works, see Appendix D.5)
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 D.6). In the
absence of layernorm, our forward pass has the form

$\displaystyle\bm{h}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)=\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)+\frac{\beta_{0}}{L^{\alpha_{L}}}\text{MHSA}\left(\bm{h}^{\ell}(\bm{x},t)\right)_{\mathfrak{s}}$
(21)

where the MHSA layer is

$\displaystyle\text{MHSA}\left(\bm{h}^{\ell}(\bm{x},t)\right)_{\mathfrak{s}}=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}\in[\mathcal{H}]}\sum_{\mathfrak{s}^{\prime}\in[\mathcal{S}]}\bm{o}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\sigma^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\bm{o}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{O\mathfrak{h}}(t)\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
(22)
$\displaystyle\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{V\mathfrak{h}}(t){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle\sigma^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\sigma\left(\bm{\mathcal{A}}^{\ell}_{\mathfrak{h}}(\bm{x},t)\right)_{\mathfrak{s}\mathfrak{s}^{\prime}}\
,\
\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\frac{1}{N^{\alpha_{\mathcal{A}}}}\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)\cdot\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t)$
$\displaystyle\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{K\mathfrak{h}}(t){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle\bm{q}_{\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{Q\mathfrak{h}}(t){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\
\ ,\ $ (23)

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

$\displaystyle\bm{g}^{\ell}_{\mathfrak{s}}(\bm{x},t)\equiv\gamma_{0}N\mathcal{H}\frac{\partial
f(\bm{x},t)}{\partial\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)}$ (24)

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

$\displaystyle{M}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t)\equiv\frac{1}{N\sqrt{\mathcal{H}}}\
{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{o}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\dot{\sigma}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}(\bm{x},t)\equiv\frac{\partial\sigma^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)}{\partial\mathcal{A}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x},t)}$
(25)

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

$\displaystyle\bm{W}^{\ell}_{O\mathfrak{h}}(t)=\bm{W}^{\ell}_{O\mathfrak{h}}(0)+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}}\Delta(\bm{x},t^{\prime})\tilde{\bm{g}}^{\ell}_{\mathfrak{s}}(\bm{x},t^{\prime})\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t^{\prime})^{\top}$
$\displaystyle\bm{W}^{\ell}_{V\mathfrak{h}}(t)=\bm{W}^{\ell}_{V\mathfrak{h}}(0)+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}\mathfrak{s}^{\prime}}\Delta(\bm{x},t^{\prime})\sigma^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}\tilde{\bm{g}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t^{\prime})\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x},t^{\prime})^{\top}$
$\displaystyle\bm{W}^{\ell}_{K\mathfrak{h}}(t)=\bm{W}^{\ell}_{K\mathfrak{h}}(0)+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\Delta(\bm{x},t^{\prime})M_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t^{\prime})\dot{\sigma}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x},t^{\prime}){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t^{\prime})^{\top}$
$\displaystyle\bm{W}^{\ell}_{Q\mathfrak{h}}(t)=\bm{W}^{\ell}_{Q\mathfrak{h}}(0)+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\Delta(\bm{x},t^{\prime})M_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t^{\prime})\dot{\sigma}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t^{\prime}){\bm{h}}^{\ell}_{\mathfrak{s}^{\prime\prime}}(\bm{x},t^{\prime})^{\top}$

We can now express the residual stream as

$\displaystyle{\bm{h}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle=\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)+\beta_{0}L^{-\alpha_{L}}\
\bar{\bm{\chi}}_{O\mathfrak{s}}^{\ell+1}(\bm{x},t)$
$\displaystyle+\eta_{0}\gamma_{0}\beta_{0}^{2}L^{-1}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}^{\prime}\sim\mathfrak{B}_{t^{\prime}}}\Delta(\bm{x}^{\prime},t^{\prime})\sum_{\mathfrak{s}^{\prime}}{\bm{g}}^{\ell+1}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){V}^{\ell\sigma}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,s)$
(26)

where we introduced the fields

$\displaystyle\bar{\bm{\chi}}_{O\mathfrak{s}}^{\ell}(\bm{x},t)=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\mathfrak{s}^{\prime}}\bm{\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\sigma^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)\
,\
\bm{\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{O\mathfrak{h}}(0)\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$

and the kernel

$\displaystyle{V}^{\ell\sigma}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,s)=\frac{1}{\mathcal{H}N}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{v}^{\ell\sigma}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})=\frac{1}{\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}\sigma_{\mathfrak{s}\mathfrak{s}^{\prime\prime}}^{\ell}\sigma^{\ell}_{\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime\prime}}V^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})$
(27)

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

$\displaystyle\bm{\chi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{V\mathfrak{h}}(0)\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
(28)

We similarly find the following for the key dynamics

$\displaystyle\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\bm{\chi}_{K\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)$
$\displaystyle+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}N^{1-\alpha_{\mathcal{A}}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}\Delta(\bm{x},t^{\prime})M_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x},t^{\prime})\dot{\sigma}^{\ell}_{\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime\prime\prime}}^{\ell}(\bm{x},t^{\prime}){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})$
$\displaystyle\bm{\chi}_{K\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{K\mathfrak{h}}(0){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
(29)

and an analogous update equation holds for the query dynamics $\bm{q}$

$\displaystyle\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\bm{\chi}_{Q\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)$
$\displaystyle+\frac{\beta_{0}\eta_{0}\gamma_{0}}{L^{1-\alpha_{L}}N^{1-\alpha_{\mathcal{A}}}}\sum_{t^{\prime}<t}\mathbb{E}_{\bm{x}^{\prime}\sim\mathfrak{B}_{t^{\prime}}}\sum_{\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\Delta(\bm{x}^{\prime},t^{\prime})M_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x}^{\prime},t^{\prime})\dot{\sigma}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}(\bm{x}^{\prime},t^{\prime})\bm{k}_{\mathfrak{h}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x}^{\prime},t^{\prime}){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})$
$\displaystyle\bm{\chi}_{Q\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)=\frac{1}{N^{\frac{3}{2}-\alpha_{\mathcal{A}}}\sqrt{\mathcal{H}}}\bm{W}^{\ell}_{Q\mathfrak{h}}(0){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
(30)

We see that if the residual stream is frozen so that
$\bm{\chi}_{K},\bm{\chi}_{Q}$ are static, the keys and queries will only
evolve in the $N,L\to\infty$ limits if $\alpha_{L}=\alpha_{\mathcal{A}}=1$ as
we argued in the main text. From these equations we can deduce the pre-
attention values
$\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)$.

Next, we examine the dynamics of the
$M^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$ variables which are defined as

$\displaystyle
M^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle=\frac{1}{N\sqrt{\mathcal{H}}}{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{o}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t).$

We can verify that $M^{\ell}_{\mathfrak{h}}$ are all
$\Theta_{N,\mathcal{H}}(1)$ throughout training by expanding the dynamics of
the attention output
$\bm{o}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$.

####  D.1.1 Backward Pass

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

##### MHSA Layer

We will start by differentiating through the MHSA layer

$\displaystyle\bm{g}^{\ell}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle=\sum_{\mathfrak{s}^{\prime}}\left(\frac{\partial\tilde{\bm{h}}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x},t)}{\partial\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)^{\top}}\right)^{\top}{\bm{g}}^{\ell+1}_{\mathfrak{s}^{\prime}}$
$\displaystyle=$ $\displaystyle\
{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)+\frac{\beta_{0}}{L^{\alpha_{L}}}\sum_{\mathfrak{s}^{\prime}}\left(\frac{\partial}{\partial\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)^{\top}}\text{MHSA}(\bm{h}^{\ell}(\bm{x},t))_{\mathfrak{s}^{\prime}}\right)^{\top}{\bm{g}}^{\ell+1}_{\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle=$ $\displaystyle\
{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)+\frac{\beta_{0}}{L^{\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{\mathfrak{s}^{\prime}}\sum_{\mathfrak{h}}\bm{W}^{\ell}_{V\mathfrak{h}}(t)^{\top}{\bm{g}}^{\ell\sigma}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
$\displaystyle+\frac{\beta_{0}}{L^{\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{\mathfrak{h}}\bm{W}^{\ell}_{Q\mathfrak{h}}(t)^{\top}\bm{k}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\frac{\beta_{0}}{L^{\alpha_{L}}\sqrt{N\mathcal{H}}}\sum_{\mathfrak{h}}\bm{W}^{\ell}_{K\mathfrak{h}}(t)^{\top}\bm{q}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$ (31) $\displaystyle=$
$\displaystyle\
{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)+\frac{\beta_{0}}{L^{\alpha_{L}}}\left[\bar{\bm{\xi}}^{\ell}_{Q\mathfrak{s}}(\bm{x},t)+\bar{\bm{\xi}}^{\ell}_{K\mathfrak{s}}(\bm{x},t)+\bar{\bm{\xi}}^{\ell}_{V\mathfrak{s}}(\bm{x},t)\right]$
$\displaystyle+\frac{\beta_{0}^{2}\eta_{0}\gamma_{0}}{L}\sum_{t<t^{\prime}}\mathbb{E}_{\bm{x}^{\prime}\sim
t^{\prime}}\Delta(\bm{x}^{\prime},t^{\prime})G^{\ell\sigma}_{O\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
$\displaystyle+\frac{\beta_{0}^{2}\eta_{0}\gamma_{0}}{L}\sum_{t<t^{\prime}}\mathbb{E}_{\bm{x}^{\prime}\sim
t^{\prime}}\Delta(\bm{x}^{\prime},t^{\prime})\left[K^{\ell
M\dot{\sigma}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+Q^{\ell
M\dot{\sigma}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(32)

where we introduced the variables

$\displaystyle\bm{g}^{\ell\sigma}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\sum_{\mathfrak{s}^{\prime}}\sigma^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)\bm{g}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\bm{k}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\sum_{\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}M^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}(\bm{x},t)\dot{\sigma}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime\prime}\mathfrak{s}}(\bm{x},t)\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\bm{q}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\sum_{\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}M^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime\prime}}(\bm{x},t)\dot{\sigma}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}(\bm{x},t)\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime\prime\prime}}^{\ell}(\bm{x},t)$

and their associated kernels

$\displaystyle
G^{\ell\sigma}_{O\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\bm{g}^{\ell\sigma}_{O\mathfrak{h}}(\bm{x},t)\cdot\bm{g}^{\ell\sigma}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle K^{\ell
M\dot{\sigma}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\bm{k}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{k}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
$\displaystyle Q^{\ell
M\dot{\sigma}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\frac{1}{N\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\bm{q}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{q}^{\ell
M\dot{\sigma}}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(33)

where we introduced the following random fields

$\displaystyle\bar{\bm{\xi}}^{\ell}_{V\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}\mathfrak{s}^{\prime}}\bm{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\sigma^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x},t)$
$\displaystyle\bm{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{V\mathfrak{h}}(0)^{\top}\bm{g}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\bar{\bm{\xi}}^{\ell}_{Q\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}\bm{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\dot{\sigma}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}\mathfrak{s}}(\bm{x},t)M_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x},t)$
$\displaystyle\bm{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{V\mathfrak{h}}(0)^{\top}\bm{g}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)$
$\displaystyle\bar{\bm{\xi}}^{\ell}_{K\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}\bm{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime\prime\prime}}(\bm{x},t)\dot{\sigma}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime\prime}\mathfrak{s}^{\prime\prime\prime}}(\bm{x},t)M_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime\prime}}^{\ell}(\bm{x},t)$
$\displaystyle\bm{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime\prime\prime}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{K\mathfrak{h}}(0)^{\top}\bm{q}_{\mathfrak{h}\mathfrak{s}^{\prime\prime\prime}}^{\ell}(\bm{x},t)$
(34)

####  D.1.2 Why we need $\alpha_{\mathcal{A}}=1$ for gradient stability

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

##### Isolating all Dependence on Initial Conditions

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

The following variables are necessary to characterize the dynamics of the
$N\mathcal{H}$-dimensional residual stream

$\displaystyle\bm{\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{O\mathfrak{h}}(0)\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\
,\
\bm{\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{O\mathfrak{h}}(0)^{\top}\bm{g}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)$
$\displaystyle\bm{\chi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{V\mathfrak{h}}(0)\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\
,\
\bm{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{V\mathfrak{h}}(0)^{\top}\bm{g}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
$\displaystyle\bm{\chi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{Q\mathfrak{h}}(0)\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\
,\
\bm{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{Q\mathfrak{h}}(0)^{\top}\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
$\displaystyle\bm{\chi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N\mathcal{H}}}\bm{W}^{\ell}_{K\mathfrak{h}}(0)\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\
,\
\bm{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)=\frac{1}{\sqrt{N}}\bm{W}^{\ell}_{K\mathfrak{h}}(0)^{\top}\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
(35)

From these variables, we can reconstruct the entire dynamics for the
$\bm{h},\bm{g}$ fields. We therefore study the moment generating functional
for the above primitive random variables. Let
$\bm{\theta}_{0}=\\{\bm{W}^{\ell}_{O\mathfrak{h}},\bm{W}^{\ell}_{V\mathfrak{h}}(0),\bm{W}^{\ell}_{K\mathfrak{h}}(0),\bm{W}^{\ell}_{Q\mathfrak{h}}(0)\\}$

$\displaystyle Z[\\{\bm{j}\\}]=\mathbb{E}_{\bm{\theta}_{0}}$
$\displaystyle\exp\left(\sum_{\ell\mathfrak{h}\mathfrak{s}t}\int
d\bm{x}\left[\bm{j}^{\chi_{O}^{\ell}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\chi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\bm{j}^{\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\xi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right]\right)$
$\displaystyle\times\exp\left(\sum_{\ell\mathfrak{h}\mathfrak{s}t}\int
d\bm{x}\left[\bm{j}^{\chi_{K}^{\ell}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\chi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\bm{j}^{\xi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right]\right)$
$\displaystyle\times\exp\left(\sum_{\ell\mathfrak{h}\mathfrak{s}t}\int
d\bm{x}\left[\bm{j}^{\chi_{Q}^{\ell}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\chi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\bm{j}^{\xi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right]\right)$
$\displaystyle\times\exp\left(\sum_{\ell\mathfrak{h}\mathfrak{s}t}\int
d\bm{x}\left[\bm{j}^{\chi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\chi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\bm{j}^{\xi^{\ell}_{V}}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right]\right)$
(36)

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 $\bm{\chi}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)$ and
its corresponding pair $\bm{\xi}^{\ell}_{\mathfrak{s}}(\bm{x},t)$

##### Attention Output Matrices

In this section we integrate over
$\bm{W}^{\ell}_{O\mathfrak{h}},\bm{W}^{\ell}_{K\mathfrak{h}},\bm{W}^{\ell}_{Q\mathfrak{h}}$.

$\displaystyle\ln\mathbb{E}_{\bm{W}^{\ell}_{O\mathfrak{h}}(0)}\exp\left(-\frac{i}{\sqrt{N}}\sum_{t\mathfrak{s}}\int
d\bm{x}\
\text{Tr}\bm{W}^{\ell}_{O\mathfrak{h}}(0)^{\top}\left[\hat{{\bm{\chi}}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)^{\top}+\frac{1}{\sqrt{\mathcal{H}}}{\bm{g}}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)\hat{\bm{\xi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)^{\top}\right]\right)$
$\displaystyle=-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{{\bm{\chi}}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{{\bm{\chi}}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})V^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})$
$\displaystyle=-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{\bm{\xi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){G}^{\ell+1}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})$
$\displaystyle-i\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{\bm{\xi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{v}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)R^{g^{\ell+1},\chi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
(37)

where we introduced the response function

$\displaystyle
R^{g^{\ell+1},\chi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\equiv-\frac{i}{N\sqrt{\mathcal{H}}}\bm{g}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\chi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}).$
(38)

##### Value Matrices

Next, we average over the value matrices $\bm{W}^{\ell}_{V\mathfrak{h}}(0)$
which gives

$\displaystyle\ln\mathbb{E}_{\bm{W}^{\ell}_{V\mathfrak{h}}(0)}\exp\left(-\frac{i}{\sqrt{N}}\sum_{t\mathfrak{s}}\int
d\bm{x}\
\text{Tr}\bm{W}^{\ell}_{V\mathfrak{h}}(0)^{\top}\left[\frac{1}{\sqrt{\mathcal{H}}}\hat{{\bm{\chi}}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)^{\top}+{\bm{g}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)^{\top}\right]\right)$
$\displaystyle=-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{{\bm{\chi}}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\hat{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){G}^{\ell}_{O\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle-i\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot{\bm{g}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})R^{h^{\ell},\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle
R^{h^{\ell},\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=-\frac{i}{N\sqrt{\mathcal{H}}}\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(39)

##### Key/Query Matrices

Next, we need to perform the averages involving $\bm{W}_{K\mathfrak{h}}(0)$
which has entries resulting in

$\displaystyle\ln\mathbb{E}_{\bm{W}^{\ell}_{K\mathfrak{h}}(0)}\exp\left(-\frac{i}{\sqrt{N}}\sum_{t\mathfrak{s}}\int
d\bm{x}\
\text{Tr}\bm{W}^{\ell}_{K\mathfrak{h}}(0)^{\top}\left[\frac{1}{\sqrt{\mathcal{H}}}\hat{{\bm{\chi}}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)^{\top}+\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)^{\top}\right]\right)$
$\displaystyle=-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{{\bm{\chi}}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\hat{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})Q^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle-i\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})R^{{h}^{\ell},\xi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle
R^{{h}^{\ell},\xi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\equiv-\frac{i}{N\sqrt{\mathcal{H}}}{\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(40)

We follow an identical procedure for the query matrices
$\bm{W}_{Q}^{\ell}(0)$.

$\displaystyle\ln\mathbb{E}_{\bm{W}^{\ell}_{Q\mathfrak{h}}(0)}\exp\left(-\frac{i}{\sqrt{N}}\sum_{t\mathfrak{s}}\int
d\bm{x}\
\text{Tr}\bm{W}^{\ell}_{Q\mathfrak{h}}(0)^{\top}\left[\frac{1}{\sqrt{\mathcal{H}}}\hat{{\bm{\chi}}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)^{\top}+\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)^{\top}\right]\right)$
$\displaystyle=-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{{\bm{\chi}}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\hat{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})K^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle-i\sum_{\mathfrak{h}tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{{\bm{\chi}}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})R^{{h}^{\ell},\xi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle
R^{{h}^{\ell},\xi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\equiv-\frac{i}{N\sqrt{\mathcal{H}}}{\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
(41)

##### Enforce Kernel Definitions

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

$\displaystyle 1=\int$
$\displaystyle\frac{dH^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})d\hat{H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})}{2\pi
iN^{-1}\mathcal{H}^{-1}}$
$\displaystyle\exp\left(\hat{H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\left[N\mathcal{H}H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})-\bm{h}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dG^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})d\hat{G}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})}{2\pi
iN^{-1}\mathcal{H}^{-1}}$
$\displaystyle\exp\left(\hat{G}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\left[N\mathcal{H}G^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})-\bm{g}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{g}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
(42)

This is repeated for all of the response functions which involve sums over
$N\mathcal{H}$ variables

$\displaystyle 1=\int$
$\displaystyle\frac{dR^{g^{\ell+1},\chi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})dR^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(-R^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)\left[NR^{g^{\ell+1},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\frac{i}{\sqrt{\mathcal{H}}}\bm{g}^{\ell+1}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\chi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dR^{h^{\ell}\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})dR^{g^{\ell}_{O}\chi^{\ell}_{V}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(-R^{g^{\ell}_{O}\chi^{\ell}_{V}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)\left[NR^{h^{\ell}\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\frac{i}{\sqrt{\mathcal{H}}}{\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dR^{h^{\ell}\xi_{K}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})dR^{q^{\ell}\chi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(-R^{q^{\ell}\chi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)\left[NR^{h^{\ell}\xi_{K}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\frac{i}{\sqrt{\mathcal{H}}}{\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dR^{h^{\ell}\xi_{Q}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})dR^{k^{\ell}\chi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(-R^{k^{\ell}\chi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)\left[NR^{h^{\ell}\xi_{Q}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\frac{i}{\sqrt{\mathcal{H}}}{\bm{h}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\hat{\bm{\xi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$

There are other kernels which are only relevant within a single head including
$\\{Q_{\mathfrak{h}}^{\ell},K_{\mathfrak{h}}^{\ell},A_{\mathfrak{h}}^{\ell},M_{\mathfrak{h}}^{\ell}\\}$.
These

$\displaystyle 1=\int$
$\displaystyle\frac{dQ^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})d\hat{Q}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(\hat{Q}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\left[NQ^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})-\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dK^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})d\hat{K}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(\hat{K}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\left[NK^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})-\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)\cdot\bm{k}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{dM^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)d\hat{M}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(\hat{M}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)\left[NM^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)-\frac{1}{\sqrt{\mathcal{H}}}\tilde{\bm{g}}^{\ell}_{\mathfrak{s}}(\bm{x},t)\cdot\bm{o}_{\mathfrak{h}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t)\right]\right)$
$\displaystyle 1=\int$
$\displaystyle\frac{d\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)d\hat{\mathcal{A}}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)}{2\pi
iN^{-1}}$
$\displaystyle\exp\left(\hat{\mathcal{A}}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)\left[N{\mathcal{A}}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)-\bm{k}_{\mathfrak{h}\mathfrak{s}}^{\ell}(\bm{x},t)\cdot\bm{q}^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right]\right)$

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

$\displaystyle\bm{Q}=\text{Vec}$
$\displaystyle\\{H^{\ell},G^{\ell},V^{\ell},K^{\ell},Q^{\ell}\\}$
$\displaystyle\cup\\{\hat{H}^{\ell},\hat{G}^{\ell},\hat{V}^{\ell},\hat{K}^{\ell},\hat{Q}^{\ell}\\}$
$\displaystyle\cup\\{R^{v^{\ell}\xi^{\ell}_{O}},R^{g^{\ell+1},\chi^{\ell}_{O}},R^{k^{\ell}\chi^{\ell}_{Q}},R^{q^{\ell},\chi^{\ell}_{K}}\\}$
$\displaystyle\cup\\{R^{h^{\ell},\xi_{V}^{\ell}},R^{h^{\ell},\xi_{K}^{\ell}},R^{h^{\ell},\xi_{Q}^{\ell}},R^{h^{\ell},\xi_{V}^{\ell}}\\}.$
(43)

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

$\displaystyle Z\propto\int$ $\displaystyle d\bm{Q}\ \exp\left(N\mathcal{H}L\
S(\bm{Q})\right)$ (44)

where the DMFT action $S(\bm{Q})$ has the form

$\displaystyle S$
$\displaystyle=\frac{1}{L}\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}[H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+G^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{G}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})]$
$\displaystyle+\frac{1}{\mathcal{H}L}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{Q}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime}){Q}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+\hat{K}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime}){K}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle+\frac{1}{\mathcal{H}L}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[\hat{V}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime}){V}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle+\frac{1}{\mathcal{H}L}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\sum_{t\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}\left[\hat{\mathcal{A}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t){\mathcal{A}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t)+\hat{M}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t){M}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}^{\ell}(\bm{x},t)\right]$
$\displaystyle-\frac{1}{\mathcal{H}L}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[R^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)R^{g^{\ell+1},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+R^{g^{\ell}_{O}\chi^{\ell}_{V}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)R^{h^{\ell}\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle-\frac{1}{\mathcal{H}L}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\left[R^{q^{\ell}\chi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)R^{h^{\ell}\xi_{K}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})+R^{k^{\ell}\chi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}^{\prime}\mathfrak{s}}(\bm{x}^{\prime},\bm{x},t^{\prime},t)R^{h^{\ell}\xi_{Q}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right]$
$\displaystyle+\frac{1}{L}\ln\mathcal{Z}_{\text{res}}+\frac{1}{L\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\sum_{\ell=1}^{L}\ln\mathcal{Z}_{\text{MHSA},\mathfrak{h}}^{\ell}$
(45)

The residual stream single site moment generating functional has the form

$\displaystyle\mathcal{Z}_{\text{res}}=\int\prod_{\ell\mathfrak{h}\mathfrak{s}\bm{x}t}$
$\displaystyle\frac{d\hat{\chi}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\chi_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\xi}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\xi_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\xi}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\xi_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\xi}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\xi_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}$
$\displaystyle\exp\left(-\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})h^{\ell}_{\mathfrak{s}}(\bm{x},t)h^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\sum_{\ell=1}^{L}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{G}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})g^{\ell}_{\mathfrak{s}}(\bm{x},t)g^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\sum_{\ell=1}^{L}\sum_{\mathfrak{h}}\sum_{t\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}\hat{M}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)g^{\ell+1}_{\mathfrak{s}}(\bm{x},t)o^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{{\chi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{{\chi}}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})V^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){G}^{\ell}_{O\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){Q}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){K}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(i\sum_{\ell\mathfrak{h}}\sum_{t\mathfrak{s}}\int
d\bm{x}\hat{\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t){\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\hat{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t){\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right)$
$\displaystyle\exp\left(i\sum_{\ell\mathfrak{h}}\sum_{t\mathfrak{s}}\int
d\bm{x}\hat{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\hat{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right)$
$\displaystyle\exp\left(-\frac{i}{\sqrt{\mathcal{H}}}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{\chi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)g^{\ell+1}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{i}{\sqrt{\mathcal{H}}}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{g^{\ell}_{O}\chi^{\ell}_{V}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{\xi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t){h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{i}{\sqrt{\mathcal{H}}}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{q^{\ell}\chi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{\xi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{i}{\sqrt{\mathcal{H}}}\sum_{\ell\mathfrak{h}}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{k^{\ell}\chi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{\xi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){h}^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
(46)

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

$\displaystyle\chi^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
$\displaystyle=u^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\frac{1}{\sqrt{\mathcal{H}}}\sum_{t^{\prime}\mathfrak{s}^{\prime}}\int
d\bm{x}^{\prime}R^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})g^{\ell+1}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})$
$\displaystyle u^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)$
$\displaystyle\sim\mathcal{GP}\left(0,V^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
(47)

We note that $h$ only depends on
$\bar{\chi}^{\ell}_{O}=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}\chi^{\ell}_{O\mathfrak{h}}\sigma_{\mathfrak{h}}^{\ell}$
so it has the form

$\displaystyle\bar{\chi}^{\ell}_{O}$
$\displaystyle=\frac{1}{\sqrt{\mathcal{H}}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}u^{\ell}_{O\mathfrak{h}}\sigma^{\ell}_{\mathfrak{h}}+\bar{R}^{v^{\ell}\xi^{\ell}_{O}}g^{\ell}\
\ ,\ \
\bar{R}^{v^{\ell}\xi^{\ell}_{O}}=\frac{1}{\mathcal{H}}\sum_{\mathfrak{h}=1}^{\mathcal{H}}R^{v^{\ell},\xi^{\ell}_{O}}_{\mathfrak{h}}\sigma^{\ell}_{\mathfrak{h}}$
(48)

This fact will be important when we take the large head limit with $N$ fixed
in Appendix D.3.

Next, we analyze the single site MGF for the hidden MHSA layers

$\displaystyle\mathcal{Z}_{\text{MHSA},\mathfrak{h}}^{\ell}=\int\prod_{\ell\mathfrak{h}\mathfrak{s}\bm{x}t}$
$\displaystyle\frac{d\hat{\xi}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\xi_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\chi}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\chi_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\chi}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\chi_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}\frac{d\hat{\chi}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)d\chi_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{2\pi}$
$\displaystyle\exp\left(-\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{Q}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})q^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{K}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)k^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{V}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})v^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)v^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-\sum_{t\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}\hat{\mathcal{A}}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime}){G}^{\ell+1}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{{\chi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{{\chi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{{\chi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{{\chi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(-\frac{1}{2}\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}\hat{{\chi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)\hat{{\chi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){H}^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\right)$
$\displaystyle\exp\left(i\sum_{t\mathfrak{s}}\int d\bm{x}\
\hat{\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t){\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\hat{\chi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t){\chi}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right)$
$\displaystyle\exp\left(i\sum_{t\mathfrak{s}}\int d\bm{x}\
\hat{\chi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t){\chi}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)+\hat{\chi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t){\chi}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)\right)$
$\displaystyle\exp\left(-i\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{g^{\ell+1},\chi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{\xi}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)v^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-i\sum_{tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{h^{\ell},\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{{\chi}}^{\ell}_{V\mathfrak{h}\mathfrak{s}}(\bm{x},t){g}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-i\sum_{\mathfrak{h}tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{{h}^{\ell},\xi^{\ell}_{Q}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{{\chi}}^{\ell}_{Q\mathfrak{h}\mathfrak{s}}(\bm{x},t)k^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
$\displaystyle\exp\left(-i\sum_{\mathfrak{h}tt^{\prime}\mathfrak{s}\mathfrak{s}^{\prime}}\int
d\bm{x}d\bm{x}^{\prime}R^{{h}^{\ell},\xi^{\ell}_{K}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})\hat{{\chi}}^{\ell}_{K\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right)$
(49)

###  D.2 Infinite $N$ (Key/Query dimension) Limit

First, we take the $N\to\infty$ limit with $\mathcal{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 $\bm{Q}$ take on their values at
the saddle point $\bm{Q}^{\star}$.

$\displaystyle\frac{\partial S(\bm{Q})}{\partial\bm{Q}}|_{\bm{Q}^{\star}}=0$
(50)

##### Saddle Point Values for Order Parameters

Under this saddle point, all of the order parameters presented will
concentrate and all neurons will become statistically independent. The
governing equations for the order parameters $\bm{Q}^{\star}$ are given in
terms of averages $\left<\cdot\right>$ over the single site densities defined
by the moment generating functionals $\mathcal{Z}$ and have the form

$\displaystyle
H^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<h^{\ell}_{\mathfrak{s}}(\bm{x},t)h^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>\
,\
G^{\ell}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<g^{\ell}_{\mathfrak{s}}(\bm{x},t)g^{\ell}_{\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>$
$\displaystyle
M^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\left<g^{\ell+1}_{\mathfrak{s}}(\bm{x},t)o^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right>\
,\
V^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<v^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)v^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>$
$\displaystyle
Q^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<q^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>\
,\
K^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)k^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})\right>$
$\displaystyle\mathcal{A}^{\ell}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},t)=\left<k^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)q^{\ell}_{\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x},t)\right>$
$\displaystyle
R^{g^{\ell+1},\chi^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\sqrt{\mathcal{H}}\left<\frac{\delta
g^{\ell+1}_{\mathfrak{s}}(\bm{x},t)}{\delta
u^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle
R^{h^{\ell},\xi_{V}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\sqrt{\mathcal{H}}\left<\frac{\delta
h^{\ell}_{\mathfrak{s}}(\bm{x},t)}{\delta
r^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle
R^{h^{\ell},\xi_{K}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\sqrt{\mathcal{H}}\left<\frac{\delta
h^{\ell}_{\mathfrak{s}}(\bm{x},t)}{\delta
r^{\ell}_{K\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle
R^{h^{\ell},\xi_{Q}^{\ell}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\sqrt{\mathcal{H}}\left<\frac{\delta
h^{\ell}_{\mathfrak{s}}(\bm{x},t)}{\delta{r}^{\ell}_{Q\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle
R^{v^{\ell}\hat{\xi}^{\ell}_{O}}_{\mathfrak{h}\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<\frac{\delta
v^{\ell}_{\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{\delta{r}^{\ell}_{O\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle
R^{g^{\ell}_{O}\chi^{\ell}_{V}}_{\mathfrak{s}\mathfrak{s}^{\prime}}(\bm{x},\bm{x}^{\prime},t,t^{\prime})=\left<\frac{\delta{g}^{\ell}_{O\mathfrak{h}\mathfrak{s}}(\bm{x},t)}{\delta{u}^{\ell}_{V\mathfrak{h}\mathfrak{s}^{\prime}}(\bm{x}^{\prime},t^{\prime})}\right>$
$\displaystyle