Theory statement
### Theory statement ###
**Definitions.** Consider a training dataset consisting of $P$
inputs $x ^{\mu}\in\mathbb{R} ^{N _0\times T}$ and associated labels
$y ^{\mu}\in\mathbb{R}$, where $\mu=1,\ldots P$. Call $X\coloneqq \textbraceleft x ^{\mu} \textbraceright _{\mu=1} ^{P}$
the set of training inputs and $Y\in\mathbb{R} ^{P}$ the vector of
training labels with $\mu$-th component $y ^{\mu}$. Consider a network
defined by Eqs. (1-4) and in particular call $f ^*$ the network
output (Eq. 4) corresponding to a test input $x ^* \in\mathbb{R} ^{N _0\times T}$.
We remind the reader of the following network hyperparameters: the
input's embedding dimension $N _0$, the hidden layer's width $N$,
the number of tokens $T$, the number of attention heads per layer
$H$, and the number of attention layers $L$.
**Assumptions.** Assume the network weights $\Theta\coloneqq\left(V ^{\left(0\right)},\textbraceleft V ^{\left(\ell\right)h}\textbraceright _{\ell,h=1} ^{L,H},a\right)$
are distributed according to the Bayesian posterior distribution defined
in Eq. 8, with temperature $\mathcal{T}>0$, while the query and key
weights $\textbraceleft W _{Q} ^{\left(\ell\right)h},W _{K} ^{\left(\ell\right)h}\textbraceright _{\ell,h=1} ^{L,H}$
are fixed.
Assume $N,N _0,P\to\infty$, with $P/N\coloneqq\alpha\in\mathbb{R} ^{+}$
and $P/(N _0H ^{L})\coloneqq\alpha _0\in\mathbb{R} ^{+}$, where $\alpha$,
$\alpha _0$ as well as other size parameters $T,H,L\in\mathbb{\mathbb{N}}$
are finite.
**Claim.** Under the above assumptions,
(1) the mean predictor under the posterior distribution (Eq. 8) is
given by
$$\mathbb{E}\left[f ^{*}\right]=k ^{\top}\cdot\left(K+\mathcal{T}\mathbb{I}\right) ^{-1}Y, \qquad (9) $$
where the average is w.r.t. to the posterior distribution (Eq. 8).
The vector $k\in\mathbb{R} ^{P\times1}$ and the matrix $K\in\mathbb{R} ^{P\times P}$
are defined in terms of a kernel function $\mathcal{K}:\mathbb{R} ^{N _0\times T}\times\mathbb{R} ^{N _0\times T}\to\mathbb{R}$
as $k ^{\mu}\coloneqq\mathcal{K}\left(x ^{*},x ^{\mu}|U\right)$ and
$K ^{\mu\nu}\coloneqq\mathcal{K}\left(x ^{\mu},x ^{\nu}|U\right)$, for
$\mu,\nu=1,\dots,P$. The kernel function is given by
\begin{equation}
\mathcal{K}\left(x,x'|U\right)=\frac{1}{H ^{L}}\sum _{\pi,\pi'\in\Pi}U ^{\pi\pi'}C _{\pi\pi'}\qquad\mathrm{with}\qquad C _{\pi\pi'}\coloneqq\frac{1}{N _0}\xi ^{\pi}\left(x\right) ^{\top}\cdot\xi ^{\pi'}\left(x'\right)\, \qquad (10)
\end{equation}
where $\xi ^{\pi}\left(x\right)$ is the ``attentioned input'' corresponding
to an input $x\in\mathbb{R} ^{N _0\times T}$, along path $\pi\in\Pi$
(Eq. 7) and $\Pi$ is the set of all attention paths for a given architecture
with $|\Pi|=H ^{L}$.
The matrix $U\in\mathbb{R} ^{H ^{L}\times H ^{L}}$, called *order
parameter*, is a positive semi-definite matrix given by
\begin{equation}
U=\underset{\tilde{U}}{\mathrm{argmin}} \ S(\tilde{U};X,Y)\, \qquad (11)
\end{equation}
where the scalar function $S$ called the *action* is defined
as
\begin{equation}
S(U;X,Y)=\mathcal{L}(U)+\alpha\mathcal{E}(U;X,Y)\, \qquad (12)
\end{equation}
The scalar function $\mathcal{\mathcal{E}}$, which we call the *energy*,
is given by
\begin{equation}
\mathcal{E}(U;X,Y)=\frac{1}{P}\ln\det\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right)+\frac{1}{P}Y ^{\top}\cdot\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right) ^{-1}\cdot Y, \qquad (13)
\end{equation}
where $K\coloneqq K(X,X|U)$ is the $P\times P$ training kernel matrix,
defined according to Eq. (10). The expression for the scalar function
$\mathcal{L}$, which we call *entropy*, is lengthy and is given
in Appendix B.1.
(2) In the particular case of a single head per layer $H=1$, $U$ is a scalar, and the entropy assumes the simple form $\mathcal{L}\left(U\right)=\sigma ^{-2\left(L+1\right)}U-\ln\left(U\right)$,
where $\sigma ^{2}$ is the variance of the Gaussian prior on the network
weights $\Theta$ (see Eq. 8).
(3) For general $H$, $\mathcal{L}\left(U\right)$ is minimized by
$U ^{\pi\pi'}=\sigma ^{2\left(L+1\right)}\delta _{\pi,\pi'}$,
which therefore is always the solution of Eq. (11) in
the GP limit defined by $\alpha\to0 ^{+}$.
(4) The matrix $U$ obeys the following relation
\begin{equation}
U ^{\pi\pi'}=\frac{1}{N}\mathbb{E}[V _{\text{eff}} ^{\pi}\cdot V _{\text{eff}} ^{\pi'\top}], \qquad (14)
\end{equation}
where $V _{\text{eff}} ^{\pi}\in\mathbb{R} ^{1\times N}$ are the effective weights along path $\pi$ (Eq. 6).
**Derivation:** See Appendix C.