Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers
Authors: Lorenzo Tiberi
Francesca Mignacco
Kazuki Irie
Haim Sompolinsky
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep learning theory
Statistical mechanics
Transformers
Kernel Methods
Gaussian Processes
Finite Width Networks
Primary area: other

Abstract: Despite the remarkable empirical performance of Transformers, their theoretical understanding remains elusive. Here, we consider a deep multi-head self-attention network, that is closely related to Transformers yet analytically tractable. We develop a statistical mechanics theory of Bayesian learning in this model, deriving exact equations for the network's predictor statistics under the finite-width thermodynamic limit, i.e., $N,P\rightarrow\infty$, $P/N=\mathcal{O}(1)$, where $N$ is the network width and $P$ is the number of training examples. Our theory shows that the predictor statistics are expressed as a sum of independent kernels, each one pairing different "attention paths", defined as information pathways through different attention heads across layers. The kernels are weighted according to a "task-relevant kernel combination" mechanism that aligns the total kernel with the task labels. As a consequence, this interplay between attention paths enhances generalization performance. Experiments confirm our findings on both synthetic and real-world sequence classification tasks. Finally, our theory explicitly relates the kernel combination mechanism to properties of the learned weights, allowing for a qualitative transfer of its insights to models trained via gradient descent. As an illustration, we demonstrate an efficient size reduction of the network, by pruning those attention heads that are deemed less relevant by our theory.

Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers
Lorenzo Tiberi 1.2
Francesca Mignacco Kazuki Irie
Haim Sompolinsky 1.255
CCntte for Brain Science, Harvard University, Cambridge, MA, USA
KKmppee Institute for the Study of Natural and Artificial Intelligence,
Harvard University, Cambridge, MA, USA
GGrauuate Center, City University of New York, NY, USA
"Joseph Henry Laboratories of Physics, Princeton University, NJ, USA Edmond and Lily Safra Center for Brain Sciences,
Hebrew University of Jerusalem, Jerusalem, Israel
itiberiefas. harvard. edu, fmilgnacco@princeton edu kirie0fas. harvard. edu, posppolinskymmce Cmcc harvard. edu
Abstract
Despite the remarkable empirical performance of transformers, their theoretical understanding remains elusive. Here, we consider a deep multi-head self-attention network, that is closely related to transformers yet analytically tractable. We de- velop a statistical mechanics theory of Bayesian learning in this model, deriving exact equations for the network's predictor statistics under the finite-width thermo- dynamic limit, i.e., N, P - co, P/N = O(1), where N is the network width and P is the number training examples. Our theory shows that the predictor statistics are expressed as a sum of independent kernels, each one pairing different attention paths, defined information pathways through different attention heads across layers. The kernels are weighted according to mechanism that aligns the total kernel with the task labels. As consequence, this interplay between attention paths enhances generalization performance. Experi- ments confirm our findings on both synthetic and real-world sequence classification tasks. Finally, our theory explicitly relates the kernel combination mechanism to properties of the learned weights, allowing for qualitative transfer of its insights to models trained via gradient descent. As an illustration, we demonstrate an efficient size reduction of the network, by pruning those attention heads that are deemed less relevant by our theory]
kernel combination
Introduction
In recent years, transformer models based on multi-head self-attention layers (16] have achieved remarkable performance at natural language processing and vision tasks 779 Yet, theoretical characterizations accounting for the success of these architectures remain sparse. Two fundamental discern task-relevant
questions remain to large extent unsolved: First, interpretability-hoov structures within the learned weights? Second, generalization-wii specific aspects the trans- former architecture transformers the combination of layer-wise multi-head organization with depth. This provides network with large number of attention paths, defined as specific sequences of heads through the attention layers. Their interplay is still poorly understood by deep learning theory.
responsible for their effective learning? We posit that one important feature of
Our code is
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In most cases, theoretical characterizations of transformers' expressivity [00, inductive bias [ 12 generalization [5] and training dynamics [16 rely on simplifying assumptions on the network architecture. A characterization of attention paths is inaccessible in these models, either because attention paths are not defined the first place, as in models consisting of a single-head [19 15],
single-layer [TO] I1, or both [12 [6] 18] 20], or because the interplay between paths cannot be fully described due to constraints imposed on the learnable weights [17. A few works consider a multi-head, multi-layer architecture 241,241 but address different questions than the present study, such as expressivity, generalization bounds, or phenomenological models. Further details on these and analogous works discussed in Appendix
One characterization of the complete transformer architecture has been obtained in the Bayesian framework under the infinite-width thermodynamic limit N + oo (and infinite number of heads 25 26]. an actively studied regime in which neural networks become equivalent to Gaussian processes (GP) [27 However, the attention paths interplay is lost in this limit because the network's 's hidden weights remain statistically independent after learning. This limitation can be overcome by considering the finite-width thermodynamic limit [29, where also the number of examples P - do such that P/N + a E Rt. this regime, for example, multi-gated deep networks showcase task-relevant interplay between gates, mediated by the learned weights [33]
H 00)
In this work, we apply the statistical mechanics theory of finite-width networks to a deep multi- head self-attention model, which closely mimics the attention paths interplay in transformers, while remaining analytically tractable. Our main contributions can be summarized as follows:
We derive exact equations for the predictor statisties under Bayesian learning of the network's value weights, at fixed query and key weights.
We shed light on the interplay between attention paths by uncovering task-relevant kernel combination mechanism, emerging beyond the GP limit (a > 0). This constructs the network's 's mean predictor as optimally weighted sum of many "path-path kernels". defined as similarity matrices between pairs of attention paths, thereby improving generalization.
We provide interpretability to this mechanism, by directly relating it to the magnitude and correlations developed by the learned weights. This allows our insights to be transferred out- side the Bayesian framework, to networks trained with gradient descent. As an application, we show that a trained network can be reduced in size with minimal performance loss, by pruning those attention paths that are deemed less relevant by our theory.
We corroborate our findings on both synthetic and real-world sequence classification tasks, illustrating the two main benefits of kernel combination: task-relevant weighting and correlation of the attention paths, respectively.
2 Model
We consider transformer-like 1 architecture consisting of a linear input projection layer; L. multi- head self-attention (MHA) layers, each having H attention heads; and a linear readout layer. The network input E &NNoT is a sequence of I tokens I, E R-,, with token index t {1...., T'}. and dimension No. The input projection layer performs the transformation
(0)) NNO
Itz
(() NxN.
(I)
where N is the hidden layers' width. With the operator "." we denote matrix matrix or matrix-vector multiplication. The f-th MHA layer with index E [1....,.]] performs the transformation
(E+1) It
vlOA
(8)
vloA RRNN
(2)
VNH
where, for each head h, we define the attention matrix (2()) TxT with matrix elements
(0 = (xx1)) NNN NNVG
x.
w (OAT
w-(e)h
It
W (e)b
Q K
w (0
E GxNo,
(3)
2

---

-
Attontion heads
B- Value matrices
protiin a pell,, yeie peiu
(b)
(c)
Attention parth
M
(a)
Kuateeattion Hatotii (orian ...
kew
Kep
Imput
ADD
Oulput w/
P
Attention path SD2
Order parameter:
(d)
KFw
+
fyretes prem 1) + (rees (omp))))))))))))) +
KCp
(1)
Figure 1: Scheme of the model and theory (a) Scheme of the model in terms of attention paths. (b) The order parameter assigns to each pair paths a weight, given by the overlap between the corresponding effective weights. (c) Alignment of the kernel PCs with the vector of task labels Y, in the finite-width (FW) vs GP regimes. (d) Kernel as the weighted sum of many path-path kernlss. Tesk-relevant kernel combination occurs in the finite-width regime (FW), but not in the GP limit, in which cross-path kernels are discarded, and same-path kernels are equally weighted. The result an improved kernel-task alignment in the finite-width regime (shown in (c)), enhancing generalization.
Here Ç is the softmax function, applied along the direction of the token index s, while G is the dimension of the query-key feature space. The linear readout returns the scalar output
(L+1)
E IR 1
Here 2₁- can stand for different options for reducing the token dimension at readout, namely reading from a specific token t* or averaging over all tokens (IT. = + Σ xt). The network's ' learnable parameters L.H IT. {vion, wh w
the input projection weights V(0); the value, query and key weights
Jen= and the readout weights a.
Comparison to the Standard Transformer. The above architecture presents two main simplifi- cations w.r.t. the standard transformer. First, the network is linear in the value weights, while the standard transformer has a nonlinear feedforward block after each MHA layer. Second, in any layer the attention (Eq. (3 is always computed direct function of the bare input æ, rather than the processed input (8)) These simplifications allow to apply back- sock-propagating kernel renormal- ization (BPKR) techniques [29] 33), enabling the characterization of the network beyond the GP limit. Despite these simplifications, the insights gained by going beyond the GP limit are substan- kernel
tial: we will show that, in the finite-width regime, an important combination- emerge, accounting for considerable improvement in generalization performance.
Attention paths formulation. Note that, despite the linearization in the value weights, the network is still highly nonlinear in the input, thanks to the attention operation (Eq, 3. This can be seen by the following equivalent description of the network (Fig. [a)). We introduce the concept attention puths, by defining a path "index" == (hg, fu2,- (.)) where hı,.... hy (1...., H). which uniquely identifies each possible combination of the head indices across layers, i.c., each possible path through the attention heads. The network output can be rewritten
((fff)= v(") .TT
(5)
W TET
where l1 is the set of all possible paths, and we define the "effective weights" as Vieaf) E 1
(ell)r 1 v(LMe, y(L-1)be-1 p(2)a2 v(1),
(6)
VNI"
and the "attentioned impul"
T
Σ
Dio tat1 241t
0 (1)/s 52 (2)/z
N. 12-202-11 1t-11
L (L-1)he--1 1n
ER.
(7)
11--
3

---

In Eq. (5), the network can be seen as deep linear network applied to a nonlinearly expanded input- the attentioned input ç". Through Eq. 47), we can see that the bare input nonlinearly expanded from an No-dimensional space to an No dimmensionl space, by means of H² nonlinear operations: one for each attention path.
The goal of our theory is to understand how the network learns to combine these different attention paths, by means of the effective weights (Eq. Note that the network also has other learnable parameters: the query and key weights, which parameterize the nonlinear expansion of the input to ç. The learning of these parameters is not described by our theory. As we will see in Sec. Bl our theory characterizes the learned effective weights (Eq. 6 for given, fixed realization of the query and key weights.
3 Theory
fundamental quest of deep learning theory is to understand how deep neural networks, which are often overparameterized, manage to avoid overfitting, achieving good generalization performance One important role is played by the specific choice of network architecture, which can impose an inductive bias towards better generalizing configurations of parameters, among the many that fit the training data. To study this problem, we adopt the Bayesian framework. Given a dataset
34 35
of P example-label pairs (x" y² ") we seek to characterize the Bayesian posterior (8) or
Gibbs distribution, over the parameters + (V10), {v(OK)
L,H 1e
P If (z" -
1 2T
(+) & exp
202 1011
(8)
Here f (x", Q) is the network output (Eq. corresponding to the input I², where we emphasize its dependence on θ, |-||| is the Frobenius norm, o² is the variance of the weights' Gaussian prior (set to a = 1 throughout this paper), and T > 0 is the error variance, or Gibbs temperature (not to be confused with the number of tokens 7). Characterizing the Gibbs distribution allows to gain insight into the inductive bias imposed by the network architecture. Indeed, note that, in overparameterized networks, the Gibbs distribution for - 0+ describes the statistics of those parameter configurations that perfectly fit the training data, with a bias towards small weights induced by the Gaussian prior. These statistics depend on the choice of network architecture, which can therefore bias the distribution towards better generalizing parameter contigurations. For T > O, parameter configurations that do not achieve perfect fitting are also allowed, which can help to prevent overfitting.
Note that, as discussed at the end of Sec. - we characterize the statistics of the weights θ (the linear projection, value, and readout weights) for fixed realization of the query and key weights. The fixed query and key weights can be given, for example, by pre- pre-training the network with gradient descent, or by some fask-informed initialization. In Sec. 4.2 we will show that the insights gained by our theory on the weights Q can also be applied to the network trained with gradient descent on all of its learnable parameters, including the query and key weights.
The main theoretical result of this work is an expression for the expectation E [f (x*)) of the network's output on a new test example x², under the Gibbs distribution (Eq. (8) In Sec. [.1 below, we provide this formal result, accompanied by a sketch of its derivation and a discussion of the significance of the infinitediimensioaal-- thermodynamic- limit under which our result is derived. In Sec. 3.2
we discuss the result's 's interpretation and its insights into the network's generalization capabilities.
3.1 Statement of theoretical results
Definitions. Consider training dataset consisting of P inputs x" NoxT and associated labels y" E I. where μ = I...... Call X := { the set of training inputs and Y E Rº the vector of training labels with w-th component y Consider a network defined by Eqs. [14) and in particular call f the network output (Eq. 4) corresponding to a test input x" E xox
Assumptions. Assume the query and key weights { (wrCt (E) ((())) I.H
J Je are fixed, while all other
weights = (V [V(0M)L L.H 11 a) are distributed according to the Bayesian posterior distribu-

---

tion defined in Eq. 8 Assume the "thermodynamic limit" N, No. P - oo, with P/N = a E R and () ag R*, where a. ag as well other size parameters T, H, L E N are finite.
Result 1. The mean predictor under the posterior distribution (Eq. 8
is given by
B = kT (K +T1)
The vector IR and the matrix K E RPPP, called training kernel, defined in terms of a kernel function K. RNxT x RNx77 + R as hn = KC (x" x²)) and K = K (x* x²), for U.V P. The kernel function is given by
KC (xxxx))=
1 HL K.w'en
Σ U Cra
with
Car
S E (x) (z)
5 (x))
(10)
No
where 5 (x)) is the "attentioned input" corresponding input NOXT, along path II, as defined in Eq. 7 The kernel function depends on a positive semi-definite matrix U IR H xH² called order parameter, which is given by
U = argmin SÚÜ
with
(11)
0
The scalar function S, called the action, consists of an "entropy" term L, and an "energy" term
E(U;X,Y) 1 det (K(U; TI) Y (K(U; +T))
(12)
where K(U;X) K is the training kernel matrix. The expression for the entropy is lengthy and is given in Appendix B.1 In the special case of H 1, U is scalar, and L (U) In (U). For general H. the entropy C (U) is always maximized by Ux ( 8,,, which therefore is the solution of Eq. (11) in the GP limit defined by a - 0+.
2(1++)
Result 2. The matrix U obeys the following relation
1 U - E E.
((ff)==
(13)
where Veeffr= E R are the effective weights along path æ defined in Eq. (6
Derivation. See Appendix I
The derivation, which uses the BPKR technique 29 can be sketched follows. Computing E [f]'] under the posterior distribution (9) involves evaluating a high- nigh-dimensional integral in the weights θ. The idea
first reduce this computation into an integration over lower-dimensional 'macroscopic' variable U. Importantly, while @ becomes
as N + oo, U remains fiiteedimenssionl. The reduced integral is an expectation of the r.h.s. of Eq. (9. treated
function
of U, under the distribution p (U) x exp c - {U)}. where S is the action defined in Eq. 111). Then, this integral can be solved in the thermodynamic limit N + co, using the saddle-point method, which implies evaluating Eq. 9 at the U that minimizes the action (cf. Eq. [1). Crucially, the end result fully characterized by this low-dimensional quantity U. commonly called order parameter in physics, which has direct interpretation in terms of the network weights, given by Eq. (3).
In practice, the results obtained in the thermodynamic limit represent good approximation also for the case of large but finite N. In this regard, the scaling of other hyperparameters with N - oo is of particular importance, especially the number of training examples P. In the GP limit, one considers P finite. This is also called the infinite-widd width thermodynamic limit because in practice, for given and typically large P. it is good approximation only for very wide networks, when N P. In contrast, here we consider the finite-width limit in which P/N
E R* (which includes the GP limit for a + 0+). As can be seen from Eq. [1]. the action gains a new term for a > O, which, as we shall discuss below, is fundamental to account for the learning of an attention paths interplay. Finally, we note that our numerical experiments (Sec. 4 we will consider Bayesian networks which are overparameterized, i.c. P NoH which is the network capacity fixed query and key weights.
5

---

3.2 Results interpretation and implications for generalization capability
Eq. a commonly found expression in thermodynamic theories of Bayesian learning, relating the network's 's mean predictor to kernel regression. In particular, the theory of kernel regression (39] suggests that generalization improves when the training kernel K is well aligned with the task, meaning its largest principal components (PCs) are well aligned with the vector of training labels Y.
Our result for the transformer's is kernel (Eq. 10 enables insights into how the transformer architecture favors this kernel-task alignment (Fig. ld).. The kernel consists of the sum, weighted by the order parameter U, of many path-pant kernels Cr each computing the similarity between the attentioned input on two attention paths π and T'. A notable property of the multi-head architecture is that, despite the number of attention heads growing only linearly with the depth L, the number of attention paths grows exponentially x H Therefore, the network has at its disposal an exponentially large number of path-path kernels, which it can learn, through U, to optimally combine into total kernel with improved task alignment.
This phenomenon, which we term task-relevant kernel combination, is indeed predicted by our results Eqs. ( (2). These state that the learned U minimizes function S (Eq. [T), which, through the energy term E, favors kernel-task alignment. This can be seen by interpreting the energy term (Eq. [2] as the negative log-likelihood of the training labels Y under a centered Gaussian distribution, whose covariance matrix is the training kernel K. This negative log- likelihood can be minimized by aligning the largest PCs of the covariance (i.e. the kernel K) as much as possible with Y (Fig. (c)).
In contrast, in the GP limit a - 0+, the action S (Eq. 11 consists only of the entropy term L. which does not contain any task relevant information. Its only effect is to attract U towards the GP limit solution U=* = = m,. Note that, in this limit, the benefits of kernel combination are lost (Fig. dd) bottom line): First, out of all the path-path kernels Cass, only the same-path kernels (T =T'iT) are used, while the cross-path kernels (σ + T') are discarded; Second, all same-path kernels are weighted equally, without making use of any task-specific information. Note that this is true not only for our simplified model. but also for the full transformer architecture under its known GP limit (25). tele-relevant kernel combination can therefore only emerge beyond the GP limit, in the finite-width regime a > 0 studied in this work.
Finally, our result Eq. (13) relates the order parameter to a macroscopic measure of the network weights, allowing for a direct interpretation of the kernel combination mechanism: correlating the effective weights across paths allows the network make use of cross-path kernels, while controlling their magnitude allows to weigh the different path-path kernels in a task-relevant manner.
4 Experiments
To corroborate our theoretical results, we "train" our model (Eqs. by sampling its weights O (i.e. all weights except the fixed query and key weights) from the posterior distribution Eq. (8), using Hamiltonian Monte Carlo sampling (see Appendix(F) FrmendixFFforr. for details). We consider the following two tasks: hidden Markov chain (HMC) classification, and one-shot image classification by in-context learning. The first task is defined on a synthetic dataset. Its purpose is to have a minimal, controllable setting to illustrate the effects of relk-nelevant kernel combination. In the second task, we will proceed to show analogous effects on classic image datasets (Omniglot [40] MNIST [11, and FaShionMNIST (42), and compare these results with those obtained from the same network trained with standard gradient descent on all of its parameters (i.e. including the query and key weights).
4.1 Hidden Markov chain sequence classification
Task definition. The HMC classification task is defined as follows (Fig. 2 a)). The j-tth example in the dataset corresponds to hidden Markov chain of ... of of length 7 30. alternating between two hidden states, qf¹ E {+, }}. The probability of transition to the opposite state (= + F) is p* The f-th chain can belong one of two classes, labeled y" =1, depending on whether pd 0.3 or nd = 0.7, respectively. The input tokens are a noisy, higher dimensional representation of the hidden states. These are given by x Vggg + nh' where Us E R.Vo are two orthogonal feature vectors corresponding to the states "4", with O (1) entries, while of is a zero-mean Gaussian noise, with (nt int"n("T) 6pm e (oPP) Pp + of P. PI), where P. and P1 are the projectors along the
(Fig.2

---

a)
Hidden states
High dinennsnn
b)
t= 2
t=
paths legend:
good head
good head
Class
good (g)
V
denoising (d)
adversarial #1 (a1)
Closs
adversarial #2 (a2)
random head
random head
d)
overlap kernel PCs with task labels
C)
Classification accuracy
Overtap
GP
95
0.5
FW
(%))
90
sampled
paths CP
pain QP
0.0
palls utd. OP
10
15
20
25
30
10a
102
10a
PC rank
N
e)
Clessification accuracy, noisier data
U at 10
at N=100
U at N 1000
95
10
2.5
(96)
90
patto jnd thewyy pate
81
g-d
a2
-2.5
85
1
2
1
C1 TD
10
10
10a
Figure 2: Hidden Markov chain task. (a) Illustration of the task. (b) Schematics of the network and its attention paths. (c) Top: Classification accuracy for varying N (theory: blue crosses, joined by blue line; samples: black dots). Red lines: GP limit for network consisting of all paths (solid), the good path (dashed), and the good and denoising paths (dotted). Bottom: Matrix elements of U, for varying N. The matrix indices are labeled with the corresponding path name, according to the legend in (b). (d) Normalized overlap, or cosine similarity, between the PCs of the kernel K and the vector of task labels Y (N 10: blue; GP limit: orange). PCs are ranked by their eigenvalues, from largest to smallest. Only the first 30 PCs are shown. (e) Same as (c), but for increased o 1 = 5 and a network consisting of only the good and denoising paths.
subspace parallel or perpendicular to the plane spanned by u. and u_.. Unless specified, = = 1. The separate parameterization of the parallel (a)) ) and perpendicular (a noise strengths is motivated by their distinct effect on task performance: while the first corrupts information about the underlying hidden states, inevitably putting an upper bound the classification accuracy, the second can always be filtered out by learning appropriate weights. We use P P= 100 examples for training. We test the network performance in terms of the classification accuracy A -. En " sign((())) where the sum over number P 1000 of test examples. Additional task details are given in Appendix G.I
4.1.1 Results
We consider network of L = 2 layers and H 2 heads per layer, with readout from the first token. The network has total of 4 attention paths, schematically depicted in Fig. (b) For this synthetic task, we design the fixed query and key weights, and therefore the network's attention paths, to clearly illustrate the effects of tesk-relevant kernel combination (for details, see Appendix (.2).
We design the first head of each layer to give rise to a "good" attention path (green path) such that a 94%. Along
network consisting of this good path alone achieves a high classification accuracy, this path, the first head makes use of the Markov nature of the task by attending exclusively to nearby tokens, and only if they correspond the same hidden state +: the second head performs uniform attention, effectively counting how many times the first head detected the same-state transition
+ + +. In contrast, each layer's second head initialized randomly. This results in the three remaining paths having chance- classification accuracy A ~ 50%, when considered in isolation. However, these paths have very different effects, when combined with the good path. We term two of these paths "adversarial" (red and purple paths) because they deteriorate the network performance, while we term the remaining path "denoising" (blue path) because it can be effectively combined with the good path to improve robustness noisy data.
7

---

+11
-1
??"
b)
A in distribution
N. (fashionMNIST)
theory FW, sampled GP
FW, theory +
A99)
80.0
85
FW, sampled
1-11-1
77.5
80
10a
104
102
10
sampled
theory
gradient descent
10
10
0.2
-04
d)
kernel-task allgnment
scores and performance drop
accuracy after pruning
0.4
GP
Overlap
Score
Drop (%)
heads pruned
omniglot
fashion MNIST
0.2
0
83.8
86.6
81.8
88.2
0.0
120
(21
80.2
88.9
Head (thh
3
71.6
82.3
PCs
Figure 3: One-shot image classification task. (a) Scheme of the task. (b) Classification accuracy in the GP limit (red line) and the knite-width regime (FW) for varying N (theory: blue crosses, joined by blue line; samples: black dots). (c) Matrix elements of U. The "theory" and "sampled" Us are for N 10. The matrix indices are labeled with the path index π = (hu, h2). (d) Kernel PCs' overlap with the task, in the GP limit and in the finite-width regime for N = 10. Only the first 50 PCs are shown. (e) Head score (blue) and performance drop (red) after pruning the head, for the model trained with gradient descent. (f) Classification accuracy of the model trained with gradient descent, after pruning a growing number of heads, in order of their head score.
In Fig. C top) we show the network's classification accuracy function of the width N (blue, solid curve), compared the GP limit (red, solid line). lower N, well into the finite-width regime, we observe considerable improvement in performance with respect to the GP limit. This can be understood in terms of an improved kernel-task alignment, as shown Fig. dd).
This improved alignment is ensured by the order parameter U, plotted in Fig [|c bottom) for varying N. For N 10. well into the mnith-width regime, the order parameter clearly implements the two main benefits of kernel combination: the possibility to weigh the path-path kernels differently, and the ability to make use of the cross-path kernels. The first benefit is particularly apparent in the suppression of all kernels associated with the adversarial paths. In contrast, when N = 1000 and the order parameter is very close to its GP limit U 5 , these paths are not suppressed, causing deterioration in performance compared to that of the good path alone (red, dashed line Fig. 2.c top)). The second benefit is apparent in the strong off-diagonals of U, anti-correlating the good and denoising paths. can see that, while also in the GP limit the denoising and good paths combined (dotted, red line in Fig, 2]c, top)) have a better performance than the good path alone (dashed, red line), the performance boost is even higher in the renormalized regime, which makes use of the coss-path kernels. This additional improvement in performance becomes more apparent with noisier data. This is shown Fig. e), where we plot the classification accuracy of the network consisting of only the good and denoising paths, data with stronger perpendicular noise 01 = 5.
4.2 One-shot image classification
Task definition. The one-shot image classification task (Fig. 3(a) formulated in an in-context learning setting. The network is presented with sequence of three image-label pairs. The first

---

two images belong to two distinct classes of a categorized dataset (Omniglot, FashionMNIST or MNIST in our case). They are assigned the label "+" or "-" in no particular order. The third image is assigned the label "?", and belongs to one of the classes of the first two images. The network has to output 1 according to the label of the matching image. The sequence is fed to the network as follows. Following the idea of the vision transformer (VIT) [8], each image is divided into p patches. The patch E pp of image a E {1, 3} corresponds to the token (aa-1)p+-i for total of = 3p tokens. We encode the labels + -.? using three fixed random vectors u+, G-, U2 E RR,, which we directly add each patch (i.e., token) of the corresponding image. We also encode the token position with additive sinusoidal positional encoding D. The network is trained on the Omniglot dataset 440 while we test its classification accuracy on both in-distribution (ID) unseen classes of Omniglot, and out-of-distribution (OOD) MNShI dataset (we also report results on MNIST in Appendix H
4.2.1 Results
We consider a network of L = 2 attention layers and H = 4 heads per layer, with average pooling readout, trained on a subset of P = 600 examples from Omniglot (analogous results for a deeper network with T. 3, H = 3 are also reported in Appendix H.2.2). For the fixed query and key weights required by our Bayesian network, we use the query and key weights obtained from training the same network using gradient descent, with N 512, - 128, and P 528k (i.e., the entire training set from Omniglot). We refer to Appendix H. for further details on this process.
The plots shown in Fig. Tig-Bare analogous to those for the HMC task (Fig. D and illustrate analogous kernel combination phenomena. Fig. B)) shows the classification accuracy for varying N. Again, we observe performance gap between the width and GP regimes. Interestingly. this improve- ment in performance is preserved also OOD, on FashionMNIST. Again, Fig. (]dd) shows that the performance gap can be understood in terms of improved kernel-task alignment: PCs that are well aligned with Y are of higher rank, and have a larger overlap than in the GP limit.
The order parameter (Fig. C), "theory" and "sampled"") for N = 10 is clearly far from its GP limit, accounting for the improvement in performance observed in the finite-width regime. We observe similar kernel combination phenomena as in the HMC task, with strong off-diagonal elements, and stronger weighting of certain paths w.r.t. others. Interestingly, the block diagonal structure of the order parameter allows for simple interpretation of the interplay between paths: correlations mostly occur between paths sharing the same head h1 in the first layer, which also determines which paths overall enhanced (h1 1) suppressed (h1 = 2)..
This structure of the order parameter transfers qualitatively well also to the network trained with gradient descent. In Fig.Blc. "gradient descent") we show an empirical order parameter, obtained by computing Eq. [3 using a single realization of the network's weights trained with gradient descent. Both the order parameter's block structure and matrix element signs are qualitatively preserved in this empirical estimate. We emphasize that the network is trained with the full set of training examples (P = 528k) rather than the restricted one used for the Bayesian network (P = 600), and all learnable parameters including the query and key weights, making this qualitative agreement more relevant to potential applications. One example application is provided below.
Application: model reduction via head pruning. Our theory allows us to prune certain heads in the model trained with gradient descent (leading to model size and compute reduction). with marginal performance loss. This is achieved by using the order parameter to assign a score to each attention head, according to its contribution to kernel combination. The sowest-scoring heads are then pruned from the model. The head h at layer é is assigned the score s (EE) = Σ Ex,x" x' e ), where ( is the set of all paths passing through that head, and U is the order parameter derived from theory. Fig. e shows the score of each head, normalized by the largest one, compared against the drop in classification accuracy caused by pruning that head. Note that the network is not retrained after pruning, but only reevaluated on the test examples. We observe a performance drop qualitatively in line with the head scores. Most importantly, the two lowest scoring heads only cause marginal drop performance. In Fig. 3) we report the classification accuracy after pruning an increasing number of heads, in order of their score. Up until the first two heads (amounting to 25% of the total number of network parameters), the in-distribution classification accuracy is only marginally worsen. Interestingly, the OOD classification accuracy is even improved, possibly indicating an overspecialization of the pruned heads in solving only the in-distribution task. Appendix H.2.3 we achieve analogous size reduction of 25% on larger model with H 8 heads.

---

Agreement with sampled statistics. Finally, we note that both figures 2 c,e) and bb,c) show good agreement of our theory with the mean predictor and order parameter sampled from Eq. While our theory becomes exact in the N, P - oc limit, the agreement holds even for small N 10. In particular, in figures 3bb,)), it holds even if N < H4 (H² + 1) /2 136, the number of independent entries in the order parameter, which is supposed to be a finite quantity in our theory.
5 Conclusion and Discussion
Conclusion. We introduce a transformer- model featuring deep multi-head self-attention. amenable to theoretical characterization within the Bayesian framework. Our results unveil the important role of attention paths in accounting for transformers' remarkable performance. We demon- strate that, in scenarios involving the interplay of attention paths at finite widths, generalization consistently improves compared to the GP regime, where such interplay is absent. Our theory explains this paths interplay in terms of task-relevant kernel combination mechanism, where the network's total kernel results from the sum many kernels, specific to pairs of paths, and optimally weighted to improve generalization. This mechanism is confirmed by experiments on both synthetic and world sequence classification tasks. More broadly, our results are relevant to the theory of deep learning, as they provide example of non-scalar kernel renormalization [33 43 in widely adopted architecture such as the transformer, illustrating its importance in accounting for network performance. Non-scalar, opposed to scalar (29]. renormalization can affect the network's 's mean predictor and therefore lead to improved generalization. Our work provides a novel interpretation of its benefits in terms of optimized kemnel-task alignment.
Interpretability of our Theory. We provide interpretability to the kernel combination mechanism, by relating it to observable structures in the network weights, specifically to their magnitude and correlations. These predicted structures transfer well outside the Bayesian framework, to the weights of networks trained with gradient descent, broadening the applicability of our theory. As an example, we show that a trained network can be reduced in size with minimal performance loss, by pruning those heads that are deemed less relevant by our theory. The large size reduction achieved (25%) appears in line with observations that a few specialized heads are responsible for most of the network performance [44), and the proposal of head pruning schemes during training insights may therefore be relevant to the quest for minimalistic and resource-efficien models 46 47
45 Our theoretical
Limitations and Outlook. The above results are enabled by our theory's 's ability to go beyond the GP limit (25), as well as incorporating a multi-head, multi-layer architecture- prerequisite for the very existence of attention paths. However, various limitations could still be addressed, opening for exciting research directions. For example, attention our model is only a function of the bare input, rather than the previous layer's postactivation, as in standard transformers. In this case, the theoretical challenge would be to disentangle the learning of attention paths interplay from the learning of the attention paths themselves, since now also the attention matrix would depend on the value weights. Another limitation of our model is its linearity in the value weights. It may be possible to heuristically extend the theory to include nonlinear MLP blocks in between attention layers, by replacing the GP path-path kernels appearing in Eq. (IO) with the corresponding GP kernels for the nonlinear case-an approach which has proven successful in deep ReLU networks for certain regimes (29). Introducing nonlinearities, strong feature learning may also emerge [4-55]. Note that, instead, our theory is readily extendable to the case of linear ML. blocks, well multiple outputs, following 22 33 Here we chose a minimal setting focusing only on those renormalization phenomena specific to the transformer architecture. Indeed, the presence of multiple outputs causes the same kind of renormalization independently of the network architecture (i.c., adding two new output indices the order parameter), while deeper linear blocks would not alter the essence of attention paths interplay, only affecting details in the entropy part of the action. Extending the theory to include skip connections also seems viable. A very open challenge, instead, is to characterize the learning of the query and key weights, which relates to the more general challenge of extending the BPKR technique to nonlinear deep networks. Finally, our approach characterizes the inductive bias imposed by the network architecture on the parameter configurations that fit the training data, but not the bias imposed by a learning algorithm. It would therefore be interesting to import our theory to methods characterizing deep neural networks' training dynamics [51+53].
10

---

Acknowledgements
We acknowledge support of the Swartz Foundation, the Kempner Institute for the Study of Natural and Artificial Intelligence at Harvard University, the Office of Naval Research (ONR) grant No. 10114-23-1-20511 and the Gatsby Charitable Foundation. This research was supported in part by grant NSF PHY-2309135 to the Kavli Institute for Theoretical Physics (KITP). FM was supported by the Simons Foundation (Award Number: 1141576). We have benefitted from helpful discussions with Alexander van Meegen, Haozhe Shan and Qianyi Li.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 2017.
Long Beach, CA, USA, December
[2] Yoon Kim, Carl Denton, Luong Hoang, and Alexander M. Rush. Structured attention networks. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
[3] Ankur P. Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), pages 1249-2255 Austin, TX, USA, November 2016.
[4] Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term memory-networt for machine reading. In Proc. Conf. Empirical Methods in Natural Language Processing (EMNLP), pages 551-561, Austin, TX, USA, November 2016.
[5] Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentiv sentence embedding. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
[6] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Int. Conf. Learning Representations (ICLR), San Diego, CA, USA, May 2015.
[7] Tom B Brown et al. Language models are few-shot learners. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Int. Conf. on Learning Representations (ICLR), Virtual only, May 2021.
[9] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT), pages 4171-4186, Minncapolis, MN, USA, June 2019.
[10] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can single attention layer learn? A study through the random features lens. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA. December 2023.
[11] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John M. Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In Proc. Int. Conf. Machine Learning (ICML), Baltimore, Maryland, USA, July 2022.
[12) Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[13] Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition be- tween positional and semantic learning in a solvable model of dot-product attention. Preprint 27Xiv:2402.03902, 2024.
11

---

[14] Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Mapping of attention mechanisms a generalized potts model. Phys. Rev. Res., 6:023057, Apr 2024.
[15] Lorenzo Noci, Chuning Li, Mufan Bill Li. Bobby He. Thomas Hofmann, Chris J. Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and--idth limit. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[16] Hongkang Li, Meng Wang, Sijia Liu, and Pin- Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In Int. Conf. on Learning Representations (ICLR), Kigali, Rwanda, May 2023.
[17] Enric Adix-A, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua M Susskind. Transformers learn through gradual rank increase. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[18] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S. Du. Scan and Snap: Understanding training dynamics and token composition in 1-layer transformer. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[19] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. Preprint arXiv:2312.10794 2023.
[20] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? In Proc. Int. Conf. on Machine Learning (ICML), Vienna, Austria, July 2024.
[21] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
[22] Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In Proc. Int. Conf. on Machine Learning (ICML) Baltimore, MD, USA, July 2022.
[23] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Proc. Ini. Conf. on Machine Learning (ICML), Vienna, Austria, July 2024.
[24] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In Int. Conf. on Learning Representations (ICLR), Vienna, Austria, May 2024.
[25] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. Proc. Int. Conf. on Machine Learning (ICML), pages 4376 4386, Virtual only, July 2020.
[26] Itay Lavie, Guy Gur-Ari, and Zohar Ringel. Towards understanding inductive bias transform- ers: A view from infinity. Preprint arXiv:2402.05173, 2024.
[27] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In Int. Conf. on Learning Representations (ICI.R, Vancouver, Canada, April 2018.
[28] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In Int. Conf. on Learning Representations (ICLR), Vancouver, Canada, April 2018.
[29] Qianyi Li and Haim Sompolinsky. Statistical mechanies of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3):031059, 2021.
[30] Boris Hanin and Alexander Zlokapa. Bayesian interpolation with deep linear networks. Proc. of the National Academy of Sciences (PNAS), 120(23), May 2023.
12

---

[31] R Pacelli, S Ariosto,, M Pastore, F Ginelli, M Gherardi, and P Rotondo. A statistical mechanics framework for bayesian deep neural networks beyond the infinite-widhh limit. Nature Machine Intelligence,5(122) 5(12):1497-1507, 2023.
[32] Hugo Cui, Florent Krzakala, and Lenka Zdeborova. Bayes-optimal learning of deep random networks of extensive-widthh In Proc. Int. Conf. on Machine Learning (ICML), Honolulu, HI, USA, July 2023.
[33] Qianyi Li and Haim Sompolinsky. Globally gated deep linear networks. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, November 2022.
[34] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical variance trade-off. Proc. of the National Academy of Sciences (PNAS), 156(32):15849-155555
2019.
[35] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
[36] Tishby Naftali, Levin Esther, and Sara A Solla. Consistent inference of probabilities in layered networks: predictions and generalizations. In Proc. Int. Joint Conf. on Neural Networks (IJCNN), pages 403-409, Washington, DC, USA, June 1989.
[37] David JC MacKay, A practical Bayesian framework for backpropagation networks. Neural computation, 4(3):448-477 1992.
[38] Radford M Neal. Bayesian learning for neural networks. Springer, 1996.
[39] Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1):2914, 2021.
[40] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266) 2015.
[41] Yann LeCun, Corinna Cortes, and Christopher JC Burges. The MNIST database of handwritten digits. URL http://yaammlecun.com/eeeee
1998.
[42] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for benchmarking machine learning algorithms. Preprini arXiv: /71.77788.77777 2017.
[43] Zechen Zhang and Haim Sompolinsky. Robust learning in bayesian parallel branching graph neural networks: The narrow width limit. Preprint arXiv:2407.18800 2024.
[44] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi- head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proc. Association for Computational Linguisties (ACL), pages 5797-5808, Florence, Italy, July 2019.
[45] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang. Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient LLMs at inference time. In Proc. Int. Conf. Machine Learning (ICML), Honolulu, Hawaii, USA, July 2023.
[46] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. Preprint arXiv:2302.13399 2023.
[47] Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? Preprint arXiv:2305.0775) 2023.
[48] Mario Geiger, Stefano Spigler, Arthur Jacot, and Mathhiu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020 113011, 2020.
13

---

[49] Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluc- tuations in mean field neural networks. In Proc. Advances Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[50] Alexander van Meegen and Haim Sompolinsky, Coding schemes in neural networks learning classification tasks. Preprint arXiv:2407.188077 2024.
[51] Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proc. of the National Academy of Sciences (PNAS), 116 (23): 1537-11546, 2019.
[52] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2023(11):
1 114009, 2023.
[53] Ychonatan Avidan, Qianyi Li, and Haim Sompolinsky. Connecting NTK and NNGP: A unified theoretical framework for neural network learning dynamics in the kernel regime. Preprint arXiv:2309.0444 04522,, 2023.
[54] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of Learning Representations (ICLR), Virtual only, April 2020.
functions? In Int. Conf. on
[55] Samy Jelassi, Michael E. Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA. November 2022.
[56) Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
[57] Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. In Int. Conf. on Learning Representations (ICLR), Vienna, Austria, 2024.
[58] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In Proc. Int. Conf. on Machine Learning (ICML), Honolulu, HI, USA, July 2023.
[59] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research (JMLR), 25(4)): 2024.
[60] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization. Preprint arXiv:2305. /9420, 2023.
[61] Mikhail Vasil' evich Fedoryuk. The saddle-point method. Nauka, Moscow, 1977.
[62] Mikhail Vasil' evich Fedoryuk. Asymptotic methods in analysis. In Analysis I: integral repre- sentations and asymptotic methods, pages 1911 Springer, 1989.
[63] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Int. Conf. on Learning Representations (ICLR), San Diego, CA, USA, May 2015.
[64] Michael Betancourt. A conceptual introduction to Hamiltonian Monte Carlo. Preprint arXiv: 171.v2270 2018
[65] Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and acceler- ated probabilistic programming in NumPyro. Preprint 1rX2v:1912.11554 2019.
[66] Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research (JMLR), 20(28):
1-6, 2019.
14

---

[67] Oriol Vinyals, Charles Blundell, Tim Lillicrap. Koray Kavukcuoglu, and Daan Wierstra. Match- ing networks for one shot learning. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 3630-3638, Barcelona, Spain, December 2016.
[68] Tristan Deleu, Tobias Würfl, Mandana Samiei, Joseph Paul Cohen, and Yoshua Bengio. Torch- meta: A meta-learning library for Py Torch. Preprint arXiv:1999006576, 2019.
15

---

Appendices
Further discussion on related theory works
16
II Theory
18
Definitions
18
A.1 Variables and parameters
18
A.2 Model definition).
18
Posterior distribution
19
A.4
Attention paths notation
19
Results enunciation
20
Predictor statistics
B.2 Order parameter interpretation
Derivation of the predictor statistics
22
Partition function
22
Weights integration
22
Integration of the auxiliary variable
26
D Derivation of the order parameter interpretation
27
Partition function
28
D.2 Integration of the input projection weights
28
D.3 Integration of the auxiliary variable
30
Experiments
32
Numerical evaluation of the order parameter
32
Hamiltonian Monte Carlo sampling
32
Hidden Markov chain classification task
32
G. Task detinition
32
G.2 Query and Key weights mitiillization
33
One-shot image classification task
35
H.I Training with gradient descent
35
H2
Additional results
35
16

---

Part I
Further discussion on related theory works
The theoretical properties of attention- based models have been investigated from various perspectives
in recent years. Different lines of works have studied the expressivity [10 22 21 54). the inductive bias 122 and the training dynamics [7 16 18 of attention layers. In particular, [54] shows that multi-layer attention networks are universal approximators for certain classes of functions, such as equivariant sequence- -o-eeennee functions, while [2]] highlights their computational limitations, demonstrating that self-attention cannot model periodic finite-state languages, nor hierarchical structure, unless the number of layers heads increases with input length. [16
22
derive error
bounds for non-linear models, with trainable queries and keys. As our work, 19
55
focus on
the training of value matrices. In (99 T0] query and key matrices are fixed and untrainable, while
[55] sets them equal to the identity. Several theoretical studies have focused on in-context learning [56-60]. However, the works mentioned above derive generalization bounds and do not provide tight characterization of the learning curves, A first tight analysis was done in [14, considering single-layer model with factored self-attention. The authors of [3] provide a tight analysis of single-layer attention model with trainable tied queries and keys and value weights fixed to the identity. Previous theoretical works have mainly focused on a single transformer block, where attention paths [4] finds that one transformer block can learn different linguistic tasks 558 shows that transformer block can learn to encode topical models. [55] shows that one transformer block can encode patch associations. A
are not defined. For instance,
according the position of the self-attention layer.
different line of works has studied attention layers in the infinite-wiee limit 25 26]. particular, [25] establishes equivalence between Gaussian processes and infinitely-wide multi-layer attention models with infinitely many attention heads. [26] leverages this framework and studies the inductive bias of infinitely-wide transformers towards permutation symmetric functions in sequence space.
17

---

Part II
Theory
A
Definitions
We recall the definitions for transformer model and theory, Sec. 3f the main text.
A.1 Variables and parameters
The hyperparameters are
T: number of tokens
L number of attention layers
number of attention heads per layer
No
input width
N
hidden layers width
G
query-key internal dimension
P:
number of training examples
T
Gibbs temperature (theory only)
a:
Gaussian prior variance (theory only)
The network weights are
v E I NxNo (ith NRN
input projection weights
value weights
a RERINN (hih (0) ERGNN
readout weights
query and key weights (fixed in the theory)
where h = 1, ...., H. and I....,I. L.
The network input is
I2 NoXT,
input sequence
It Ro:
single token in the sequence
where I....TT.
We append an additional index 1 = 1..... to quantities that refer to a specific example in training set, such as, e.g.
x RR NoXT
j-tth example the training set p-th training label
y" E
We further define
X = {I#) pmm
Y = f3 mn 1=1
the vector of training labels
A.2
Model definition
The network performs the following sequence of input transformations (Fig. The input projection layer is defined as
(())
(14)
N
18

---

Attention path rt: n4 (1) - 02 (1)
h (2)
heaad
head
head
No
linear projection (0)
NJ
head
head
heaad (23
(2)
output
heaa
hec.dd (F+)
input
Iakers
(H)
(2)
2: (C-1)
Layer
Nn
head (h)
N
VÍDA
W
W2
N
No
Figure Schematic representation the architecture under consideration.
The output from each attention layer é = l,.... L is defined recursively as:
H
T
1
v(th
x.
t t= 1=
(15)
VNH
k=1
where, for each head h, define the attention matrix (2((hh RxxT with matrix elements
1 5 (
NoVG T T WCAT (e)
(16)
The function ( denotes the softmax applied along the direction of the token indexed by s. The linear readout is defined
24 (L+1)
(17)
where we adopt the unified notation
[T+
T
(L+1)
average over all the tokens
(18)
readout from a single token at position t = t" in order to consider different options of token readout, adopted different tasks.
D
A.3
Posterior distribution
We are interested in computing the posterior distribution over the network's weights θ L,H (v(0) Je,h=1 .)). This is analogous to the Gibbs distribution in statistical physics and
defined as
P
p (() exp
|||l
(19)
2T
20
p=1
where || is the Frobenius norm, and Z indicates the normalization, also called partition function.
A.4 Attention paths notation
In the main text, we introduce the concept of attention paths, by defining path index
(hi,..... with hi....,h hr {1....,H] H}
(20)
19

---

which uniquely identifies each possible combination of the head indices, i.e. each possible path through the attention heads (Fig. The network output can be written as
vill v(0). g
(21)
VHINNO TEl
where II the set of possible paths, and we define the effective weights as
V
v(L)a, v(l- p(2)a, v()),,
V(ef)r ERIXN
(22)
VNL
and the attentioned input
T
s==
Σ
(1)b (2)h2 ... Di
( (L)bg 1t"
atc 2
(23)
fo
In order to present our theoretical results and their derivation, it is useful to also introduce the following notation.
Partial-pati indices. We indicate with
hhe, hL)
with he,..., hy E {1, (1...., H)
(24)
a collection head indices, from layer up to layer L. For 2or 1, T1 = v is the path index defined above, Eq. 20 and in the main text. For l + 1, ne corresponds more generally starting from layer l. We call the set of such paths II, For 1, II, = II is the collection of all complete paths defined above and in the main text.
partial path,
Attentioned input at layer l.
We define the atlentioned input at layer l, for I...., as
el ((())eeee
Σ
(4) 1olhe Tte Digten ...
12.
(L)bz
g(0) E R
(25)
were we remind the reader the compact notation defined in Eq. 18 subsequent layers are related by
The attentioned inputs at
H Σ
(
((--1)be 1 1((t-)he 1
(26)
VNH
be 1
Recall that g(t-1)he-1me = ((e-1))e according to the notation defined in Eq. 24
We also give specific definitions for the cases and f L For O, we define (00)== E*. corresponding to the attentioned input defined above, Eq. 23 and in the main text. The following relation holds
g(0)m
((0),
g(1)=
((0) E No.
(27)
VNo
For l = L + 1 we define
H ((D)2 g(D)hz
g((+1)
(28)
h
B Results enunciation
Here we enunciate our theoretical results, for which a derivation given in Appendix C and Appendix D
20

---

B.1 Predictor statistics
derive (Appendix C the statistics the network prediction (x*) = I r². These are given by
on a new test example
(f*) = k² (K + TT) (6f*)) = Ktest - T (K + + k
(29)
(30)
where (.) denotes the expectation under the posterior distribution Eq. 19 and öʃ l" (
The quantities K Keett_ E IR, k E IRP, and K RRPP are defined in terms of kernel function
K NNxT E NoXT + R. For p.v I.... P we define Kiet == K. (x*, x²), he == IC [x 2x*,
and K = (x",, x²).
We derive the following form for the kernel
2'2 = HL
Σ U Caxe (x, x')
with Car (I,x') = TxT (5) (x)) (31) No
"I''EI1
where 5" (x:) is the attentioned input corresponding to an input æ while U. called order parameter, path-by-patth matrix of size [II] x [II], where II| = H² denotes the size of the set II. The value of the order parameter determined self-coonistentlly by minimizing a scalar function S, called action in the physics literature. This is defined
L+1
= -L -- (vU+tt)
- Σ c (V0)
[(())))(t)))))) (lett)-1) + aε (UN)x.)) (32)
8=1
where, for any matrix M, we define
C (M) = -2---i (M) + In det (M)
(33)
and
E :(VW),Y) 1 In det (K(UO);) T1) + r (K(U") + 711 "
-Y (34)
Note that K (( X) is simply the kernel defined above, through Eq. 31 where we have explicitly written its dependence on (1)) and X. Here we have defined collection of matrix order parameters (e) for L = I....,,, L, of size [Hee] x [Ale| were [1te| = =++1- denotes the size of the set 1le- We called (()) U. The order parameter (L++) is instead scalar. We also defined ((++1) for é I..... L as the matrix of size [1l]] x [[1g] with elements
(f+1)m u 1hetthheeeee = (((+1))4 Sne. ext
where Te, Th II, and T4+1, Tét1 E l1411 are the partial-path indices defined in Eq. 24 hy, hp = 1,..., H.
while
The action must be minimized w.r.t. all of the order parameters. Note that in the main text we defined as function of (T() = alone. By this we mean evaluated at its minimum for a fixed (1(1) so that only the minimization w.r.t. (1) left to perform.
The cels-consitt solution for the order parameters is obtained numerically by minimizing Eq. 32 with gradient descent methods. Details on the procedure are given Appendix
B.2 Order parameter interpretation
We derive (Appendix D) the following expression for the order parameter in terms of the network weights
Una
1 N
(((lll)x
((elt)xT
n, n| II.
(35)
where (+) denotes statistical averaging over the posterior distribution Eq. 19
21

---

C Derivation of the predictor statistics
Here we provide the derivation of our result for the predictor statistics, reported in Sec. B.1
C.1 Partition function
Let us recall the posterior distribution (Eq. 19
1 20
p (+)
Z exp
())
y²1²
|0|
(36)
27 1=1
where Z indicates the normalization, also called partition function.
is useful to introduce a vector t E R²' P auxiliary variables in order to linearize the squared error in Eq. 36 The partition function then reads
x / ex
Σ 1tth (f" (f" 8²)) vr))
(37)
where is the imaginary unit, tth is the j-thh component of t, while ft is shorthand for f (r" θ) and we use the symbol "x" indicate that are neglecting multiplicative constants. We also defined
[|t]l") and DO de cxp (- [|0]())
the Gaussian integration measures Dt dt exp
Calling f pP+1 add source term it ++1 p 1 to the exponential in the partition function, corresponding to the test example x² x 1- obtaining
P+1
z [tt DDDO exp
tt f"
(38)
#=1
p=1
In what follows we write [+1 as Z. without writing its argument explicitly. The partition function Eq. alloww obtain the predictor statistics the test example by differentiation
(f*) = (fPt1) = d d In (Z)
(39)
1P+1=
d
(85*) = (55 d d(WP+1) In (Z) 1
(40)
where 8f - and (-) denotes statistical averaging over the posterior distribution (Eq, [9p. In what follows, we will neglect any multiplicative constants in the partition function, that are independent of itP+1, since they are irrelevant for the computation of Eq. 39 and Eq. 40
(pP+))
C.2 Weights integration
We proceed to compute the partition function (Eq, 38 by integrating all of the network weights O and finally the auxiliary variable t. We use the back propagatinn kernel renormalization (BPKR) method 22 33), which consists in the successive integration of the weights, starting from the last layer down to the first layer.
Let us write Eq. 38 as
2 Dlexp -1
> H1 T())
(41)
where we define
P+1 ex
(t)
fM Σ
nt
(42)
Below we focus on computing T()). what follows, we will write 7 (t) as 1, without writing its argument explicitly.
22

---

C.2.1 Integration over the readout weights
We start by integrating over the readout weights In Eq. 42 obtaining
substitute f ((4tt)),
P+1
== /
DO exp
exp
8
gl
glL+))jug
(43)
In what follows, it is useful to define q (I++1)
Σ
P+1 m=1
tug(t.+
Dy and general
P+1 Vee(t)mes
(0me
Wy E II
(44)
We compute the Gaussian integral in the readout weights a, obtaining
xx / DeLL exp (-2 1 N g(lt1)T gla+t))
(45)
E.H Sem--)
With ( we indicate the collection of weights (
i.e. all weights
up to layer L.
C.2.2 Integration over the weights at layer Plugging g (L+1)
VNH Emt=1
V(l)ee
(L)MM into Eq.
we have
1
H Σ
I x De exp
(hhee T(LMMT, ((
(46)
N NH
hi, =1
We perform the integral the value weights {V(M) product N identical integrals of the form
H Jn=1 by noticing that it is given by the
/
II
N
II
H
dv
exp
H
(L)hr n
Ashi.ing V(L)),
(47)
20
f,'=1
hu =1
where we defined
d² o²
(L)he (1)9% dz
Amn, = Se +
N NHH
qi
(48)
with i, i' = I,..., N and hr, M², = L,... H. Here q. (L)b,, indicates the i-th element of the vector (L)hee E N We may consider as NH NH matrix, whose two indices the pairs (1, and My). With this notation, the result the value weights integration
I x / (Dll-1) exp (- - In det (A)
(49)
2
Using the matrix determinant lemma
we find
Nn (1+ I x / DolL-D) exp N-
p(ltt)))
(50)
where we defined the scalar R(++1)
02 0²
H
R(L+1)
Σ
q (Lia . (LiNe
(51)
NH
11=1
det (1+ (1+99]) a where g is vector.
23

---

We now enforce the identity of Eq. 51 by Fourier representation of the Dirac delta function, introducing the auxiliary scalar variable (1++)) We obtain
I x / (o exp
+
1 1N
U((+1)
(.0t)R((+ R((+1))
1
N
Nin In (1 R((ttl))
2 2
2
1 o² U(L+1 NH
H
Σ
(52)
In the statistical physics language, R(L+1) and U(L+1) are called order parameters. Since N - oc, we can solve the integral in R(++1) with the saddle point method [6 62 The value ((++1) at the saddle point is
R(++1) -00(4))-1111
(53)
Therefore, obtain
I / av DeLL-)) cxp
N (viltt)
1 o² NNH
ext U g( g(ori (54)
NL Ellz
where we defined the "entropy" term
C (VlL+t) = U(++1) In (UIL+t)
(55)
introduced the [TT,.| [[T.]] matrix (L++)) with matrix elements
1lL+ U(L+1)8
(56)
ext
Here TL, m2 E II are the partial-path indices defined in Eq. 24 which in this case coincide with hy, ht = I,..., H. While the definition of (L(++) may appear superfluous here, it is useful to perform the proof by induction in Sec. C.2.4
C.2.3 Integration over the weights at layer
Next, perform the integration over the weights layer L 1. The steps almost identical those taken in Sec. C.2.2, but will lead to the introduction of a matrix of order parameters. After this layer, we will be able to provide the results for the integration of the weights at subsequent layers by induction.
Plugging (L)=z= /NH
ΣH E gl (1-11)a 1)b2-192 192 into Eq. 54 get
x / (U(L+I)) exp
2 ((((Itt)))
H
Σ
Σ
U ((L+1)m2.t x
2
hL Ae-1ihy_1
FL ElT
Xg (b-1))x.-18TT
v(l-1)me-IT V(L-I)7's (L--2)1 10/ (57)
24

---

Once again that the integral in the value weights {V(L-1)h) H identical integrals of the form
s= is given by the product of N
I 11-
dV
1)hg-1
X
H
exp
1
vLh-1)be-- A
10'
(58)
20
f,1'
11
hL
where we defined
Aih -i.
0² o²
N
N
Σ
q:
(1-1)) u (
ext
(59)
Elll
with i, 1,....,NN and h2, 2 i....,HH. H.
Again, we may consider an NH NH matrix, whose two indices run over the pairs (i, hL-1) and (7', 1/1)). With this notation, the result of the value weights integration is
/ aU(Lt)
N (Filtt)) N In det det()))
I
exp
(60)
Using the matrix determinant lemma we find
I x /
WUl++ exp
N (IilLtt)- -
N In det (I U((+t),
R(w))
(61)
where we introduced the [TT.| x [TT.. matrix R(L),, with matrix elements defined as
0² o² N NH hz-1=1
((L)TL.TT.
Σ qle-1)he- g(b-1)he-ITt,
TL, nT. E II
(62)
With the same procedure as in Sec. C.2.2
we introduce the order parameter R(L) and its conjugate
ULL), also a |11z| x [I, matrix. We have
x f
N (viltt))
exp
2
+ Tr (v(L)
1N
det
(T (T+ (-(t+t1, R(LL) ext
2 2
1 r²
H
Σ
Σ qll-l))ee T (b-
(63)
NH
ALI Elle
Elle
h
Again, we solve the integral in RLL) with the saddle point method. The value of R(2) the saddle point is
R(L) 02T(()-1 - (((+1)-11
ext
(64)
Therefore we obtain
/
N (IiLLT)) + N No (1(L))
exp
1 0²
1- (1-1))1)--))1111111 . glu-N)rimn) (65)
extt
2 N H²
Th-1,7 Ellz-1
det (1+Q. QT) = det (1+ U Q Q). with Q and U being x and matrices respectively.
25

---

where we give a more general definition of the entropy Eq. 55 such that it can take a matrix argument
= Tr (U(L).ee u( == 0²
((l+t))11 + In det (Uit)-eetttt)
(66)
ext
ext
ext
and we introduced the |11 x |11--1 1| matrix Uea with matrix elements ((()ma-nmte U Ulllhe 1i ext
((L))L 1
(67)
ext
Here TL-1, L-1 E 1- and TL. 2 II and are the partial-path indices defined in Eq. 24 while h 1/21 = ,....., H.
C.2.4 Integration over the weights at generic layer
We can now compute the integration over the remaining value weights by induction. We claim that, after integration of the weights at layer é, I will have the form
I x / aull) (Q(t))
N C (ilLt))
exp
2
N +
(U(l)_((tt)))
N
(vietn).
2
exst
+
2
extt
((00
(68)
2NH
ext
he=1
Here we defined collection of matrix order parameters, one for each integrated layer. The order parameter (1() is a partial-path-by ppatt-b--partiall-----
matrix of size [11 x [[l)) were [1ie] = H 4+1-1
denotes the size of the set II. We also defined U(e+1) the matrix of size [[Tg] [Te]] with elements
(((+))me. U metteiieeeet = (((-1))44 Sngnng
(69)
ext
where Te, T/. 11, and Te+1.TE41 E l1++1 are the partial-pati patial-paahhhh indices defined in Eq. 24 while he, he 1,..., H..
Eq. 68 verified for layer = L 1, which we derived in Sec. C.2.3
The induction step, integrating
over the weights at layer e 1, is done by plugging (f)= >n
(f--1)he-1 ql 1)he 1
into Eq. 68 and applying exactly the same steps presented in Sec. C.2.3
C.3
Integration of the auxiliary variable t.
After integrating all of the network weights, we have
P / Dtexp ( -1 Σ yhu (tt)
Z x
(70)
with
I(t) & /
L+1 II t=1
(
N (DiLHt)
N Σ L (v10)) + 2 t=1
exp
2
(1) (0) (0) } (71)
2 N
plugging (0)= EP-ll (10) into Eq.
71 we see that we need to perform the following
integral in / E ERRP
exp
(
1 [ P+ Kurat P+1
I
dlu
exp
-
1 (K + TIl) ,1 1 -P 1 +1)) ((P+) T.))
2
(72)
26

---

where for convenience report here the kernel definitions given in Sec. B.1 The quantities Kiest E E. k RR, and E RPx are defined in terms of kernel function K: Nox 7 x E Nox - R. For μ4, l........ form of the kernel
...... P we define Kust == K. (a* x²), ktt = K. (x" x()), and Ku = K. ,
x²). The
1C (xxxx))
H Σ
U Capz (2,
(73)
mil
with
Cain (0, = g
((0)m g(0)) (a²)
(74)
No
Computing the Gaussian integral (Eq.
72) we obtain
2 x
/
L+1 nl
avt))
exp (-
No ((U00)_)))
exp
1 P+
Krest - KT (K+ k) pt p+ (K + +T)
(75)
2
where we report here for convenience the definition of the action given in Sec. B.1
s (1000) ((100)())
1+11) ) -c -- (v(itt)
Σ c (v
(2ltt)-1) aE (U(1))
ext
(76)
1=1
with
1 E (v
1 Y p
Indet (K TI)
(K + +T) Y
(77)
In the limit N, P + 0c, P a E solve the integrals in {T(O)
with the saddle point
method 61 62). The partition function therefore takes the final form Je=1
N
7 x exp
|Kust RT (K+ - ks nt P+1 at P (K +T1))
(78)
where we recall that Krest, and K all depend on U11), which must be evaluated the minimum of the action Eq. 76 with respect to all of its arguments {17(0
L+1
Tm-1
Differentiating by wP+1 the partition function Eq. 78 (see Eq. 39 and Eq, 40) we obtain the results
for the predictor mean (Eq. 29 and variance (Eq. 30 presented in Sec.
D Derivation of the order parameter interpretation
Here provide the derivation of our result the order parameter interpretation, exposed in Sec. B.2 The derivation is almost identical to that for the predictor statistics given in Sec. C What follows below should be considered continuation of Sec. which we refer for definitions.
For convenience, we report here the result we want to derive
(((eff)m h (eff))]]T
(79)
where TT. 7/ E II are the path indices defined in Eq. 24 while (-) denotes statistical averaging over the posterior distribution Eq. 19 We also recall the deniiiion of the network effective weights
V a. ((Lhe, p(b-1)he-1... p(2)k, V(MM,
(ifef) E RIXN
(80)
27

---

D.1 Partition function
As in Sec. C.I we start from the partition function
P
2 x / DIDO
Σ
2thh (ftt (ft
(81)
To the partition function, we add a source term x Ex V(oft) D0 with q
(1)m
E R
DrDO exp (- Σ y Σ
Z /
Vieif)m )
tt fm
(82)
V VHH
such that differenniiaingg by q+
(1)=
allows us obtain
N (v(nttt v(ott)=T) -H 1
dz
(83)
7
dy dq
(0) (0)
-0
where 9a.i indicates the i-th component of the vector q- matrix whose T1-th component 9- (0)==
(0) R. while qo RN the
As in Sec. [ we proceed to compute the partition function (Eq. 82 by integrating all of the network
weights O, and finally the auxiliary variable t. We make the following observation. In Eq. 82
can write explicitly f" VNHL Emen,, (eff)m (())T.,) Furthermore, in Sec. C.2.1
define q (L++1) E1= (ag(L+)) and in general
(0ml Σ
me E IT
(84)
Then Eq. 82 takes the form
Z α / ex
-1
Σ
thy +
Σ wielh
((fff)
(g + q (m))
(85)
VNHL
that the steps for computing Eq. 85 are identical to those in Sec. C until the integration over the input projection weights (1)),
Renaming (1)= + q: (1)m
+ (1))
D.2 Integration of the input projection weights
After integrating all of the network weights except the input projection ((0) we have
/ Dtexp
(-> th yd I(t)
Z
(86)
with
L+1 / DV /
|
av
N (vLL-N)
N
I(I)
exp
ext
t=2
o²
U
(2)m
9 (1)x1 ((()))) (87)
2 NHE
m. ElL,
We now substitute (m) + (1)) + qo obtaining
(1)= in Eq. 87 as well as (1)= - Σ V(0) ()))
NNo
28

---

((1) x / DV(o) /
L+1
II
(8)
N L (
N L
Σ
exp
+ 2
+
2
1=
t=
o²
U (2) (1)xTT (1)m1
2 2
ext
.9.
o²
Σ
a y (0)TT V(OTT v(0) (0)
HL
mi.
NVNolL
Σ Fi-
(22)m (2)mT T v(0) g O)
ext
9o
(88)
The integral in ((0) in Eq.
88 has the form
N No
N
No
N
No
dv 19
exp
20
ij
V 18'
Jig
v ij
(89)
=1j=1
i=1
j.j'=1
i=1 j=1
where we defined
Ajs'
6355
0²
a²
Σ
(2))m (0)x1 (0)=
(90)
No
Ti
Ell,
0²
Jij N.NoHL
Σ
() (1)m (0)m)
Lofi 4j
(91)
where = I,...,,,,
.N, while =.j'=11 1,.... No. and with the notation Le,t
(1)m and q:
(0)=
we indicate the
i-th component of the vectors qf" (1)m and g (0)= respectively.
We may consider Aj.y' the elements of No x No matrix A. With this notation, the partition function after performing the integral Eq. 89 is
Z x / Dtexp
Etng
yuy
/
L+1 II
du
1=2
x exp
N
C (U(ltt))
+ Σ
L
c
(1(0) ext
N In det (A)
2 (xiT
2
2
F=2
1 o2
(1)=1
m22)
m2 (0)m (0)g
((2)0
2 N
g-
ext
N Ho
exxt
qj
13
exx
Eill
(92)
We now differentiate the partition function in Eq.
92 by gr specified by Eq.
83 obtaining
(y(effmm. -
f Dtexp I the yu
L+1 II
deico
7 741.0
(0) -00-0) 0)
o²
ext
2
N No T.
d2
Σ
(12) (0)ps
exct
qj
A J5. q3
(2)oi exct
x
Pr Ell,
N
(U(L+t))
+ ((10).
N L
N In det (A) (93)
X exp
+ 2
2
C
ext
2
2=2
29

---

The remaining steps are the same as in Sec. C We use the matrix determinant lemma and the Woodbury matrix identity Coentityyresseettttt respectively to express In det (A) and A- in terms of [I[]] [II matrix R()) with elements
o² No (1)m T g(())i
Ti, . II
whose identity we enforce by Fourier representation of the Dirac delta function, introducing the auxiliary [[h] x [h matrix U11) The result of these operations is
((eft)m v(elt)aiT
(0) = 0) / Dtexp -1 -1= you I
L++ duio
II
dR
x
7 (11.
o²
(12)) extt
1| - R((). U + R(!)). (1 (1+ U12l -R(1)
U(2) ext
R()) U
x
x exp
+ !
N (
+
N
Σ
L (1(...
2
1=2
+ 2
N
Tr (U(W
2
1
N In det
[I+
((2)
exet
R!
(() (0)=) (0 000 } (94)
2 HHH
a1 Ell
As in Sec. C we solve the integral in R1()) with the saddle point method. The value of R((2) at the saddle point
R -- ((2)-1
(95)
Plugging this back into Eq. 94 we obtain
(p(oftmm ploffrrit
/ Dlexp
I
L+1
II auvi)
-1
Z q : =0)
umm
UW)sss exp
C c (UULtt))
N
+
N
Σ
C
(10)) (11000ett))))
2
2
E=1
o²
Σ
1( ( ((()
(96)
NHL
mI Ell
D.3 Integration of the auxiliary variable t
The calculation of the integral in t follows that in Sec.
C33
We obtain
((elf) p(elf)a
(
(97)
N
vas
where (v00)p(00)) denotes the expectation under the distribution
p (UTa) x /
[++ II aU
N
L+1 1-1
exp
(98)
det (I+ (1+Q.U.Q) = det (I U.Q QT Q) with Q and U being m x n and n x matrices respectively.
(U+Q-U-QT) = 1 [++ U Q Q) U Q with Q and U being m x n and n x
matrices respectively.
30

---

where the action S is the same defined in Sec. C.3 Eq. 76 Exactly as in Sec. expectation using the saddle point method. under the limit N, P oc. N arrive at the final result
C.3
we compute the We therefore
(v ((eff) U(
(99)
where (1)) is the value the saddle point of S.
31

---

Part III
Experiments
E
Numerical evaluation of the order parameter
The order parameter in theory is defined self-consistently the minimum of an action, We determine this order parameter numerically, by minimizing the action using the Adam optimizer [63] As we describe in Appendix [H.2.1] for each N (out of 10 choices among N NE 10 10"00,00:000 E {1, 2, 5) b E 2,,2,3) 3) we search for the optimal temperature among 10 choices T E {a 10-8, 1.0, 1.5; a E {1, 2.5, 5, 7.5] b E {1. 2) resulting in 100 total configurations. single A100-40GB GPU. The learning rate is optimized
Eq. 32
Each such run takes less than 12 hours for each configuration by sweeping among {10 10 : a {1,5, 8} E {0, 1, 2, 3}} for 10 iterations at the beginning of each run, and by selecting the one that achieves the lowest energy term (averaged over these 10 first optimization steps).
Hamiltonian Monte Carlo sampling
We sample the network weights from the posterior distribution (Eq. [9p using Hamiltonian Monte Carlo sampling 664. Specifically, we use the NumPyro implementation of the No U-Turn Sampler (NUTS) 66)
For N < 100, we run 10 independent chains, consisting of 1000 warm-up steps and 1000 sampling steps. We keep one sample every 10, for total of 1000 samples. Due to limited computational resources, the number of samples is smaller and varies for > 100, while we did not take samples at all for very large N 2 1000. Note however, that the large N regime is the least relevant to sample, since the network is approaching the GP limit. The most important validation of our theory is performed for the smaller values of N.
Each run sampling a model a given N takes less than 12 hours on a single A A100-40GB GPU. The number of runs to sample all model widths for the hidden Markov chain classification task and the one-shot image classification task is 12.
Regarding the temperature T of the Bayesian posterior Eq. 19 this is set differently depending on the task. For the HMC task, we find the temperature TT not be relevant for improving the network's classification accuracy. We therefore set it to small, but finite value of T 0.01. In contrast, for the one-shot image classification task, we find tuning the temperature be particularly important to prevent overfitting. Therefore, we always tune the temperature to the value giving the optimal classification accuracy for the given network depth N. We refer to Appendix H.2.1 for details on the temperature values and its optimization process.
G Hidden Markov chain classification task
Here we give additional details on the hidden Markov chain (HMC) classification task.
G.1 Task definition
Here we recall the task definition, providing few additional details.
The 4-th example in the dataset corresponds to hidden Markov chain of ..... of of length T 30, alternating between two hidden states, of E {+, -). The transition probability to the opposite state
+ " ")" p". In other words, describing the "4"" states as one hot vectors + 0 and
the transition probability matrix of the hidden Matkov chain is
(' pl
D
p
1 1
32

---

a)
b)
adversarial #1
adversarial #2
2=2
t=1
paths legend:
2
good head
good head
Examples
good (g)
20
denoising (d)
adversarial (a1)
40
0
adversarial #2 (a2)
random head
random head
good
denoising
0
0.4
Examples
0.6
C)
20
0.2
Classification accuracy, noisier data
0.4
40
95
X X Se
A (96)
GP
RN
IX:
patiic g+d, therry paths g+d. GP
1.5
Examples
1.0
85
20
1.0
10a
10a
10a
0.5
40
0.5
0
25
0
25
Examples
Examples
Figure 5: Hidden Markov chain task. (a) Schematics of the network and its attention paths. (b) Kernels. Same-path kernels associated with the 4 paths shown in (a), total kernel in the GP limit, and total kernel for N 10 in the renormalized regime (RN). Examples ordered by class (the first half correspond to the first class, the second half correspond to the second class). (c) Classification accuracy for network consisting of only the good and denoising paths, Tl| = and σд = 5. The figure is analogous to Fig. f), with the difference that we have replaced the random head inv in the denoising path with uniform attention head.
both the and y axes
The u-th chain can belong to one of two classes, labeled y" 11, depending on whether p = 0.3 pd 0.7 respectively.
The network is presented with visible states the input tokens which are noisy, higher dimensional representation of the hidden states. These are given by
x oi + T/
Here Ug E No are two orthogonal feature vectors corresponding to the states "4". We set No 200 and
V2 if S 100
5 V2
if i S 100
U+,
U+,
otherwise
otherwise
The term of is zero-mean Gaussian noise, with (nt" ( where Pl and P1 the projectors along the subspace parallel perpendicular the plane spanned by e4 and e..
( P1 PI+oPPI Pi).
We also preprocess the dataset in two ways: adding to each chain beginning-oo-ssetennee of-sentence (bos) token of zeros at position t = 0, and concatenating to each token a one-hot positional encoding vector of size T (i.e. the number of tokens including the bos token).
We use P 100 examples for training and p 1000 examples for testing. For this task, we find the temperature T to not be relevant for improving the network performance. We therefore set it to small, but finite value = 0.01. A finite value T is required to sample the predictor statistics, hence comparing the theoretical results with samples.
G.2 Query and Key weights initialization
Here give the details of the initialization of the fixed query and key weights.
We recall that we consider a network of T. layers and heads per layer, with readout from the first token (i.e.. t" = 1). The network has total of 4 attention paths (Fig. B(a). For the first head
33

---

of each layer, we make good choice of the fixed query and key weights, which defines a "good" attention path, achieving a good classification accuracy (ef. Sec. J.T]in the main text). The remaining heads are initialized at random. Below we give the initialization details of these good and random heads.
Let us recall here the definition of the attention matrix ( E RRT for head h at layer l. Its matrix elements are defined as
(2)h = =<(NvG'''
w-lth WIOOT
w-(Oh w-la E R Q
Zt
where 8, 1= ...... T, while Ç is the softmax function, applied along the direction of the token index 8, and the dimension the query-key feature space. We directly initialize the query-key matrix product W = w with. Note that W(OA is an [No + (T + 1)] x [No + (T + 1)] matrix, because we have appended one-hot positional encoding vector to the No-dimensional input tokens (see Sec. G.1 In order to define the heads, it is convenient to decompose W(()) into the block structure
W Wr(th pf
wole w-len
w(th = 3
where W-Ohh E (R(++)x(T++) acts only on the one-hot positional encoding subspace, (((h) NooxNNo No acts only on the subspace of the tokens' "features", and wr(t)h w-(t)AT E Nooox(TT)) mix these two subspaces. The scalar < 3 < oo is a parameter controlling the "hardness" of the softmax
4p
function (for / + oc, the softmax becomes a hardmax). We set it to 8 10 for all heads.
G.2.1
Good heads
Let us define the good heads, i.e. (((hh for h and & 1, Note that the goal here is not to define heads that are provably good at solving the task, but rather to make a good guess for their initialization, based on our knowledge of the nature nature of the task.
First layer. For the head h = 1, E = we define
(x+ - )ー (o+ -
w- --(1)1 = 0
No
and
((11
3 do, 16,,'
t I....,T T
PP
where
( ee'
is the component of w w((2)a indices t,t'.
Second layer. The head 1, l 2 implements uniform attention Q(2)1 T. o,...,T. is defined by w-(1)1 w(1)1 (1)1 w = 0, and w pp =
T+1 Vs, t
fp
As discussed in the main text, the attention path defined by the good heads achieves good classi- fication accuracy. We can give an intuition as to why this is the case. Intuitively speaking, in the limit of hardmax attention (3 xc) and noise (aI 0 - 0). the attention path is "counting" the number of times a token has remained in the same state after a new step in the Markov chain. Indeed, the first head "detects" when there has not been change of state between adjacent tokens. It does so by either attending nearby tokens if and only if they are in the same state, or attending "nothing" (in the sense of the zero beginning-of-sentenee token). Then, the second head sums over the tokens attended by the first head, thereby "counting" the number of times a token has not changed state. More generally, outside the above mentioned limit, we can say that the good attention path focuses on the two most relevant pieces of information needed to solve the task: First, it preferentially pays attention to nearby tokens, which is important because of the memoryless nature of the Markov process; Second, it is able to detect the type of transition occurrying between nearby tokens (i.e. remaining in the same state, or changing state), which is important to distinguish between the two classes, since they differ by their transition probability.
34

---

G.2.2
Random heads
Let us define the random heads, i.e. W(0) for h = 2 and f = 1.2. These are initialized with Gaussian identically and independently distributed entries
Wrlen Pp
10.1)
W PP
~V(0.1)
is
No
[
w
(0,
N (0,1)
VNo
VNo
= L...=.....A. No and Vt, = 0, O,..., T. Note that take care of proper normalization of the above matrices, depending on which subspaces they act upon (i.e. the "features" or the "one-hot positions" subspaces).
As mentioned in the main text (Sec. 4.1 p. the random heads introduce three additional paths: two adversarial paths, deteriorating the performance of the good path, and one "denoising" path, improving the good path performance. We can get an intuition same-path kernels, Fig- (b).. We can see that both the good-path and the two adversarial-path kernels appear very structured, with sharp excursions their values for different pairs of examples. However, while the good-pal path kernel structure appears to be aligned with the task, well distinguishing the two classes, the ad adversarial-p -pthh kernels structure appears random w.r.t. to the task. We can expect that adding these adversarial kernels to the good one would destroy it's task-relevant structure, can be visually understood from the total GP kernel. In contrast, the total renormalized kernel, in which the adversarial-pan kernels do not contribute, preserves the task-relevant structure. Differently, the denoising-patis kernel appears less structured and more uniform, with weaker noisy excursions. In fact, what suspect is that there is nothing special about the specific realization of the random head involved in the denoising path, which is just implementing a noisy version of uniform attention. We verify this by substituting the random head with implementing uniform attention and repeating the same experiment shown in Fig. 2 in the main text. This is shown in Fig. [ic)), were we plot the classification accuracy of the network consisting of the good and denoising paths alone, for the case of | = 1 and Tj = 5. We can see that the results are completely analogous to those shown in Fig. f in the main text.
why this so by looking at their associated
H One-shot image classification task
H.1 Training with gradient descent
Here provide details of gradient descent training of our transformer-like model (Sec. 2 for the one-shot image classification task. We use the Omniglot dataset with the standard I 1028/172/432-spp for the train/validation/tes. class splits [67] implemented in torchmeta [68). We use the Adam optimizer [63] using an initial learning rate of 3e and a batch size of 128 for 10 epochs. We check the validation accuracy every 1000 steps and select the final model as the one that achieves the best validation accuracy. We use the binary regression loss as in the theory. Unlike in the theory, here we train all the model parameters including the key and query projection weight matrices. We set N = 512. All the models considered in this work can be trained on a single A100-40GB GPU within less than 2 hours.
H.2 Additional results
Here we report further results on the one-shot image classification task.
H.2.1
Optimal temperature
For the Bayesian model, we find tuning the Gibbs temperature T to be particularly important to optimally perform the task. All results for the network's 's classification accuracy presented in Sec. 4.2 and below therefore shown at the optimal temperature for the given N, obtained by scanning the set of temperatures T E {a 10-b, 1.0, 1.5), where E 21, 5, 7.5} and b E (1,
Note that temperature is optimized only for the in-distribution classification accuracy. In particular, we do not optimize for temperature when testing distr---ittiibutionn but rather keep the optimal temperature determined by evaluating the network in-distribution.
35

---

a)
L=2 H=4
A vs N, in distribution
A vs N, ood (fashionMNIST)
A vs N. ood (MNIST)
92.0
A (%0
80.0
(96)
85
91.5
80
10a
10
10a
10
10a
10
N
N
L=3 H=3
distribution
N. ood (MNIST)
A N. (fashlonMNIST)
87.5
(%6)
80.0
(%6
85.0
(%6)
85.0
77.5
82.5
82.5
-4
10a
102
10a
10a
10a
10a
N
b)
sampled
theory
202383835
Figure 6: (a) Classification accuracy for varying N, tested both in-distribution (Omniglot) and out-of- distribution (MNIST, fashionMNIST). We report the results for the network architecture discussed in the main text (L 2, H = 4) and deeper one with one more layer (L 3. Hl = 3). Theory: blue crosses, joined by blue curve. GP limit: red line. (b) Order parameter for the L = 3, H = 3 network,
for N 10. Similarly the T. 2, H network shown in the main text, the order parameter
showcases attention paths interplay, presenting strong off aff-diagonal elements that deviate from the GP limit.
For the network considered in the main text, find the following optimal temperatures: T 0.5 for N = 10, 20, 50, 100; T = 0.25 for N = 200, 500, 1000; T = 0.1 for N 2000, 5000; T = 0.075 for N 10000. Note that the optimal temperature grows consistently N becomes smaller. This can be readily understood by inspecting the equation for the mean predictor (Eq. 29), which we report here for convenience
(f))
=T TT (K 4K)-
where we recall that
KM
HT. Σ Una Car
1- 1,..,, P
where Cre' is a path-path kernel. For decreasing N, we typically observe the order parameter growing in overall magnitude, which in turn affects the magnitude of the kernel K. As consequence, also the optimal temperature needs to be rescaled. The fact that U grows in magnitude for smaller N can be understood from the energy term (Eq. (34) in the action (Eq, 32). We discussed in Sec. that this can be seen as the negative log- likelihooo of the labels vector under centered Gaussian distribution, whose covariance matrix is the kernel K. While the most effective way to minimize the energy term that described the main text (Sec. i.e. aligning the kernel with the task, one more trivial way is to increase the log-likelihood variance in all directions (i.e. increasing the kernel's overall magnitude).
all of our experiments, we always observe this phenomenon of growing magnitude to a certain extent.
36

---

a)
Head scores and performance drop
b)
accuracy after pruning
Drop (%)
heads pruned
omniglot
fashion MNIST
MNIST
Score
0
83.8
86.6
86.6
81.8
88.2
88.2
(12
(212
(271
(2)1
2
80.2
88.9
88.9
71.6
82.3
58.8
Head (4)h
accuracy after pruning
C)
d)
heads pruned
omniglot
fashion
MNIST
Head scores performance drop
Drop (9)
0
86.0
79.8
91.9
Score
1
85.5
81.0
91.1
2
84.6
79.0
90.8
3 83.7
83.1
90.3
117
125
121
120
(2)4
1277
f218
4
82.4
77.9
90.9
Head (f)h
5
80.4
80.6
87.6
Figure 7: Heads pruning experiment. (a,c) Head score (blue) and performance drop (red) after pruning the head, for the network trained with gradient descent. (a) smaller network considered in the main text (L= 2, H = 4); (c) larger network (L = 2, H 8) (b,d) Classification accuracy of the model trained with gradient descent, after pruning growing number of heads, in order of their head score. (b) smaller network considered in the main text (L 2, H = 4) (d) larger network (L 2, H
H.2.2 Classification accuracy
We perform the same experiments on the classification accuracy for a deeper network (I. = 3, H 3). The results are shown in Fig. 6 Note that here also report the test accuracy MNIST, which was not shown the main text. can see that the results discussed in the main text are confirmed also for the deeper network. In particular, for the in-distribution classification accuracy, we consistently observe a performance improvement in the renormalized regime, with respect to the GP limit. When testing out of distribution.
see that in the best cases (fashionMNIST for (L 2, H 4); MNIST for (L = 3, H 3)), the performance improvement is preserved, or at the very worst (MNIST for (L = 2. H 4); fashionMNIST for (L 3, H 3)) the GP and renormalized regime show comparable performance.
H.2.3 Heads pruning
We repeat the head pruning experiment for network with more heads per layer (L 2, H 8). We find that we can prune with marginal performance loss a similar percentage of heads as in the smaller network. The results are shown in Fig. 7 Again, see that the head scores determined by our theory are qualitatively in line with the performance loss caused by pruning the corresponding head (Fig. c)). Note that we do expect perfect alignment between the two quantities. The most important fact to verify is that the low scoring heads correspond to a small drop in performance. In Fig. d) we show the classification accuracy after pruning an increasing number of heads, in order of their score. Up to 4 heads (25% of the network size) the in- in-distribution performance has only marginal drop, identical that obtained after pruning heads the smaller network considered in the main text (also accounting to 25% of its size, Fig. bb)). We can also see that pruning up to three heads improves the oft-of-ditrre -fit--iiiiiiutiion performance on fashionMNIST,, as we also observed for the smaller network.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We confirm that main claims in abstract and introduction accurately reflect our contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Sec. paaragrpph "Comparison to the Standard Transformer", and Sec. paragraph "Limitations and Outlook".
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
38

---

Answer: [No]
Justification: The paper uses the methods and approach of theoretical physics. As such, it provides theoretical results in the form of formulae (clearly enunciated in Sec. and
Appendix [B), alongside their exact derivations (provided in completeness in Appendices C
and D). definitios and assumptions are stated in Sections 2 and as well as Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We describe all technical details in the paper. particular, experimental details are provided in Sec. and Appendix III Our code and scripts are also provided in the supplemental material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
39

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our code provided in the supplemental material. We will make it publicly available upon acceptance. The datasets we use in our experiments are publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https https//oooSSuuiisseeeoe
/ //nipscco wwtppcpcccggpciceccccccoocccccce (elpii//giidee//oooooo
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all such technical Detailshhppe--prrameee in Appendices E and Our code and scripts are also provided in the supplemental material.
H.I
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The theory and sampling of the classification accuracy and order parameter have no error bars by definition. Our general predictions regarding the behavior these quantities at different network widths confirmed on very different tasks; Furthermore, for the second task (one-shot image classification) we confirm these predictions for three conell-ooniiuurations that are different in terms of number of layers and heads.
Guidelines:
40

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots. The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the corresponding information in Appendices E and H.I
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https://neurrri wwtp://bbb..ttiicccccuuieeeeedde
Answer: [Yes]
Justification:
We have nothing specific to report here.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We have nothing specific to report here.
41

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Parts of the code we use are originally forked from some public repositories. We explicitly acknowledge this and include the corresponding license files/text.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
42

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thooe cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The documentations of our code are provided in the readme files in the supplemental material.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We have no such experiment in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: We have no such experiment this work.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
43

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
44


Reviews:
[{"review_id": "QlI7p13jGn", "reviewer": "Reviewer_5PZg", "summary": "This work studies the mechanism of attention in deep Transformers. The theory shows that the prediction process can be expressed as a combination of kernels of different attention paths. Experiments are conducted to verify the theory, which also motivates a pruning method on different attention heads.", "questions": "1. Can you have some discussion about the following related work about model pruning? This work also finds that some heads can be pruned during the inference. However, their experiments are implemented by real-world state-of-the-art LLMs. Their results are more complete in terms of performance, efficiency, and on more challenging tasks. \n\nLiu et al., ICML 2023. Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time.", "limitations": "There is no potential negative societal impact of their work. The extension of the proposed method and analysis to larger models is the biggest concern.", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The mechanism of the attention path is interesting and reasonable. \n2. The theory and experiments support each other and make sense.", "weaknesses": "1. The experiments are simple. It will be better to verify the mechanism on larger datasets, larger models, and more complicated tasks. Experiments use Transformers of 2 or 3 layers, which are not deep. However, the abstract claims a \"deep multi-head self-attention network\".  \n\n2. The presentation is somehow weird. The order of figures in Figures 2 and 3 is strange. Also, Figures 2 and 3 can be divided into multiple figures, since the experiments in these subfigures are about different things. \n\n3. Some references on the theoretical mechanism of Transformers in learning are missing.  \n\nOlsson et al., 2022. In-context Learning and Induction Heads.\n\nLi et al., ICML 2024. How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (This work also includes model pruning based on the mechanism of Transformers)\n\nNichani et al., ICML 2024. How Transformers Learn Causal Structure with Gradient Descent.\n\nReddy et al. ICLR 2024. The mechanistic basis of data dependence and abrupt learning in an in-context classification task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.\n\nAt the moment, however, we believe that our contributions have not been evaluated in the proper context. Before reading this response, we kindly ask the reviewer to read our global response, which should help in better contextualizing our work.\n\n1. *The experiments are simple. It will be better to verify the mechanism on larger datasets, larger models, and more complicated tasks. Experiments use Transformers of 2 or 3 layers, which are not deep. However, the abstract claims a \"deep multi-head self-attention network\".*\n\nAs explained in the global response, our work is theory-focused and provides exact analytical expressions for the network predictor. In this context, experiments are typically designed to serve two purposes: 1) To convince the reader, who may not want to check the details of the theoretical derivations, that the analytical predictions are exact. 2) To provide minimal, easily interpretable examples illustrating the theoretical insights. Specifically, we focused on experiments with only 2 or 3 layers and 2 to 4 heads because the resulting order parameter remained visually interpretable in that case.\n\nIn the Bayesian framework adopted here, numerically validating our derivations on larger models and datasets would require an exponentially larger computational cost—much larger than with gradient descent—other than being technically challenging, because it requires sampling the network weights from an high dimensional posterior distribution (Eq. 8). We believe this to be out of the scope of our work: our theoretical results are rigorous and apply for any depth, under the thermodynamic limit specified by the theory.\n\nWe acknowledge that the term “deep” is used in the literature with varying meanings. In the context of theory, where analytical results have predominantly been obtained for single-layer attention, we followed the convention of using “deep” to emphasize that our derivations apply to models with an arbitrary number of layers. However, if the reviewer believes it is more appropriate, we are willing to change the term “deep” in the abstract to “multi-layer.”\n\n2. *The presentation is somehow weird. The order of figures in Figures 2 and 3 is strange. Also, Figures 2 and 3 can be divided into multiple figures, since the experiments in these subfigures are about different things.*\n\nWe thank the reviewer for spotting this problem. We have changed the order in both figures and made them more readable. The updated figures are available in the joint PDF. We have decided not to split the figures, because all the subplots are closely related and together show how the path combination mechanism affects the performance. For instance, in order to understand both the accuracy trend and the pruning mechanism in Fig. 3, it is useful to look at the structure of the order parameter. \n\n3. *Some references on the theoretical mechanism of Transformers in learning are missing.*\n\nWe thank the reviewer for pointing us towards these missing references. These will appear in the revised version as follows. Refs. [Li et al., ICML 2024; Nichani et al., ICML 2024] will be included in the Introduction, in the discussion of other theory-focused, closely related works. Refs. [Olsson et al., 2022; Reddy et al. ICLR 2024] are also highly relevant, but more suitable for the Discussion section, since they focus on empirical methods [Olsson et al., 2022] and phenomenological models [Reddy et al. ICLR 2024] that are more distant from the focus of our work on exact derivations of the network’s predictor.\n\n4. *Can you have some discussion about the following related work about model pruning? This work also finds that some heads can be pruned during the inference. However, their experiments are implemented by real-world state-of-the-art LLMs. Their results are more complete in terms of performance, efficiency, and on more challenging tasks.*\n\nWe thank the reviewer for pointing us towards this interesting work [1], which will be discussed in the Discussion section of the revised manuscript.\n\nHowever, we would like to stress that paper [1] is very different in scope and methodology to our work. Our work is theory-focused, and aims at providing exact analytical expressions that describe the network’s learning mechanisms. The cited work [1] is application-focused, and aims at providing a competitive pruning algorithm. In our work, the application to head pruning is neither the main result nor the focus, and we are not making any claim about it being competitive: we only use it to elucidate the kernel combination mechanism and its interpretation.\n\nWe believe it unfair to compare our experiments to those appearing in [1]. As we explain in the global response, our work is part of a line of research aimed at achieving a better analytical understanding of deep learning. Given the current state of theory, simplifying assumptions are typically necessary to obtain exact analytical results. Given these constraints, it is not fair to compare the standard transformer used in practice to our simplified model. Real-world applications would anyways be out of scope for our work. We direct the reviewer to our global rebuttal, in which we explain how, when seen in the right context, our work presents novel theoretical insights, as well as several advancements compared to the simplifying assumptions of other closely related theoretical works.\n\nWe hope our response clarifies the scope of this work, and highlights our true contributions on both advancing the frontiers of the theoretical research on transformer-like models, and deep learning theory in general.  In our view, the current rating is unfair in light of these contributions. If you find our response convincing/useful, please consider amending the score. Thank you very much.\n\n[1] Liu et al., ICML 2023. Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HVoAWRvZAy", "reviewer": "Reviewer_9bWs", "summary": "This paper provides an interpretation of wide (large embedding dimension, $N$) multi-head attention-only transformers solving in-context learning tasks as performing high-dimensional kernel combining. The paper derives the exact statistics of the predictor, when the number of training examples $P$, and $N \\to \\infty$, but $P/N \\to \\mathcal{O}(1) > 0$. The main result is the revelation that the predictor statistics can be expressed as the weighted sum of independent kernels, each pairing two attention paths through the multiple heads across the layers, and the weights depend on the relevance of the path to the task to be solved. Capitalising on this result, the paper proposes a network size-reduction method by efficient pruning of certain the attention paths with marginal performance loss.", "questions": "(1) I am curious about the authors' thoughts on what can happen when the attention is computed from the outputs of the previous layer and not always from the bare inputs?\n\n(2) In order to prune the attention paths, the order parameters $U_{\\pi, \\pi'}$ (Eq. (13)) must be computed, which involves averaging over the Gibbs distribution (Eq. (8)). How can this be computed in real architectures?", "limitations": "(1) The limitations are already stated towards the end of the paper; (i) it does not consider the learning of query-key matrices, which is a crucial step in functioning of the attention mechanisms (ii) the attention is always computed from the bare inputs, which limits the applicability of the theory to practical transformers. However, these do not reduce the value of the contribution of the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "(1) The paper introduces a novel way of looking at generalization capabilities of a transformer through the lens of the statistical mechanics theory in finite thermodynamic limit. \nSpecifically, the result shows that transformer-like architectures perform a task-relevant kernel combining, where each kernel is based on a path through the multi-head, multi-layer network. This is a crucial direction that offers valuable insights into the nature of the attention mechanisms.\n\n(2) The proofs are clearly written with the necessary steps (barring some steps like saddle-point methods) to follow.\n\n(3) The experimental results (Figure 3) greatly help in understanding the theoretical results (Section 3) better. \n\n(4) The contents and organisation of the paper is easy to understand.", "weaknesses": "(1) One minor weakness is that the experiments are limited to binary classification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time reviewing our work, for the many positive comments, and for the intriguing questions. \n\nBelow are our replies to the reviewer’s questions, as well as minor concerns.\n\n1. *One minor weakness is that the experiments are limited to binary classification.*\n\nWe agree that considering a network with multiple outputs would be interesting, and readily implementable within our theory. We would like to share with the reviewer why we actively decided to not include these kinds of experiments in this first work exploring our theory. The renormalization effect coming from having a number M of multiple outputs is the same for any architecture (e.g. the same as in deep linear networks [Li, Sompolinsky, PRX 2021]), and results in the order parameter gaining an additional pair of “output indices” ranging from 1 to M. In this work, we wanted to focus on illustrating exclusively the renormalization effects that are proper of the transformer architecture—-that of attention paths combination. Exploring the interplay between the two renormalization effects (multiple outputs and multiple paths) could be the topic of future work, with potentially exciting applications to language and next-token prediction tasks. One caveat in this context, is that one would need to be careful of a too large number of outputs M, which may spoil the thermodynamic limit which assumes M finite.\n\n2. *I am curious about the authors' thoughts on what can happen when the attention is computed from the outputs of the previous layer and not always from the bare inputs?*\n\nIt seems reasonable to believe that an attention paths combination mechanism would still persist in this case. The theoretical challenge, however, would be to disentangle the learning of attention paths interplay from the learning of the attention paths themselves. Indeed, when the attention is a function of the bare input, our results for the renormalized kernel \n$$K=\\sum_{\\pi,\\pi'\\in\\mathrm{\\Pi}}U^{\\pi\\pi'}C_{\\pi\\pi'}$$\nshow that the value weights’ statistics enter only in determining the order parameter $U$, that is the interplay between attention paths. If the attention were to be a function of the previous layer’s output, instead, also the attention path kernels $C$ would depend on the value weights, which in this case affect the learning of the attention paths themselves. In other words, the value weights would influence performance in two key ways: (i) by optimizing the representation of inputs to the attention heads, and (ii) by improving the recombination of attention paths.\n\nObtaining analytical progress in this scenario is challenging, primarily due to the introduction of non-linearities in the trained weights through the attention’s softmax. A first step to advance in this computation could be to consider linear attention heads, which may simplify the analysis.\n\nThis discussion will be included in the revised manuscript, as also prompted by a question from reviewer Y9xT.\n\n3. *In order to prune the attention paths, the order parameters 𝑈𝜋,𝜋′  (Eq. (13)) must be computed, which involves averaging over the Gibbs distribution (Eq. (8)). How can this be computed in real architectures?*\n\nWe would like to be sure to consider the case meant by the reviewer with “real architectures”, so let us consider a few scenarios below:\n\n- a. For our model in the Bayesian framework, the order parameter can be equivalently computed by either solving the self-consistent equation (as in figure 3d “theory”), or sampling the posterior (as in figure 3d “sampled”).\n- b. For our model trained with gradient descent, one can just take the learned query and key weights, and plug them into procedure (a) to obtain the order parameter for the value weights. Another more empirical approach that we did not try, would be to get multiple “samples” of the value weights by retraining the network with gradient descent, but for fixed query and key weights coming from the first iteration of the training (this is somehow what we did to compute the order parameter in figure 3d “gradient descent”, only that there we used a single “sample”). \n- c. For the more standard transformer with attention as a function of the bare input and trained with gradient descent, one could still do something similar to (b). The difference would be that now also the attention matrix $\\Omega$ would depend on the value weights. Instead of taking the learned query and key weights, and plugging them into procedure (a), one could just take the entire attention matrix $\\Omega$ (as computed on each example from the network's activations) and plug it into procedure (a).\n- d. If we also add nonlinear MLP blocks in between each head, one could still apply the same procedure as in (c),  but, when arriving to step (a), he could only proceed by sampling the posterior, while the option of solving the self-consistent equation would not be available due to the self-consistent equation not being defined for the nonlinear case. A possible way to define the self-consistent equation could be to heuristically extend it to the nonlinear case, using the same heuristic argument as in [Li, Sompolinsky, PRX, 2021]. There, the analytical results for a deep linear network are heuristically extended to a network with ReLU activations, by replacing the linear GP kernels with the ReLU GP kernels. Despite having a more complex renormalization, also in our case the renormalized kernel can be seen as a linear combination, weighted by the order parameter, of many “GP” kernels, i.e., the path-path kernels. Therefore, we foresee how a similar argument could potentially be applied in our case.\n\n4. *...barring some steps like saddle-point methods*\n\nThe revised manuscript will have an improved explanation of the saddle point technique. We ask the reviewer, if possible, to help us by further elaborating on the unclear points.\n\nOnce again, we thank the reviewer for this thorough review and the very positive rating.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NNxrdqTlyo", "reviewer": "Reviewer_Y9xT", "summary": "The paper investigates Bayesian learning of the value weight matrices of a deep multi-head attention network without MLPs employing the back-propagating kernel renormalization (BPKR) [1] technique in the linear regime, where the training set size $P$ scales with the width $N$, i.e., $P/N = O(1)$. It finds that the network’s kernel is a sum of constituent kernels operating on different pairs of attention paths, rescaled based on their alignment with the target task. This renormalization enhances generalization compared to previously studied infinite-width (GP) limits. The paper validates its theoretical predictions through experiments on a synthetic task and in-context classification with simple image datasets, finding qualitative agreement with gradient-based training. Additionally, it shows that the theory’s predictions can be used to prune less relevant attention heads without significantly impacting performance.\n\n[1] Li, Q. and Sompolinsky, H., 2021. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3), p.031059.", "questions": "5. If I understand correctly, unlike the dense case, thanks to the transformer architecture, the kernel is not rescaled by a scalar renormalization variable, resulting in more interesting outcomes compared to [1]. Can you elaborate more on this?\n6. In Fig. 3 (d), is the GD map obtained for the same network from which the trained query and key weights were obtained for the theory?\n7. In general, in your experiments, do you always apply attention to the input at all depths of the network?\n8. What do you mean by “efficiently” (line 57)? \n\n[1] Li, Q. and Sompolinsky, H., 2021. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3), p.031059.", "limitations": "The paper adequately lists its limitations in Section 5. I do not foresee any potential negative societal impacts arising from this study.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- **Originality:** The application of BPKR to attention models and the characterization of the network’s kernel as a task-relevant weighted sum of path-path kernels is novel.\n- **Quality:** The statistical mechanics analysis is technically sound and validated with comprehensive numerical experiments on both synthetic and real data, demonstrating the robustness of the findings. The paper provides the code to reproduce its empirical results.\n- **Clarity:** The paper is generally clear, although there are areas for improvement (see Weaknesses section).\n- **Significance:** The paper considers simplified transformer architectures, which remain challenging for theoretical approaches. The developed theory extends beyond the GP limit, predicting task-adaptivity properties that are necessary to explain the success of modern learning algorithms and are not captured by infinite-width limits. While the insights are intuitive and may not significantly enhance the understanding of transformers or their success, the theory is quantitative and offers non-trivial generalization predictions.", "weaknesses": "1. **Strong assumptions:** The analysis relies on several very strong assumptions: (i) linearity of the network output in the value weights, (ii) applying the attention at any depth on the network input, (iii) considering frozen (and already learned) query and key matrices, and (iv) a system at equilibrium (i.e., a Gibbs distribution over the parameters of the model). These assumptions clearly limit the relevance of the results.  Although some of the empirical results seem to suggest that the last two assumptions can be relaxed, the impact of (i) and (ii) on the conclusions remains unclear. It would be beneficial to discuss the implications of these assumptions further.\n2. **Background:** The paper is dense and lacks sufficient background material to aid the reader in following it and understanding where the results come from. The BPKR technique and its assumptions and rationale are not adequately introduced. Additionally, the paper mentions the \"network’s kernel\" early on without really specifying which object it is considering (as there are various kernels in deep learning, e.g., NTK etc.). Providing a more detailed introduction and some intuition on BPKR, along with its assumptions and high-level steps in the main manuscript, would significantly improve clarity.\n3. **Clarity and completeness:** Sec. 4.1.1 is not very clear and easy to follow. It would be helpful to explain the structure of the different heads within the main text, rather than repeatedly directing the reader to long appendices.\n4. **Typos:** The manuscript contains several typos (“thermodinamic” in line 106, “generalizaton” in line 147, “task-specifc” in line 161, “cathegorized” in line 241, “taks” in line 262, “on this regard” in line 330).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time reviewing our work and for many positive comments. We’d like to address the reviewer’s remaining concerns as follows.\n\n1. *The analysis relies on several very strong assumptions: (i) linearity of the network output in the value weights, (ii) applying the attention at any depth on the network input, [...] the impact of (i) and (ii) on the conclusions remains unclear. It would be beneficial to discuss the implications of these assumptions further.*\n\nWe agree with the reviewer that further discussion on assumptions (i) and (ii) would be beneficial.\n\nRegarding assumption (i), we discuss in the paper a potential way forward, following the heuristic arguments in [Li, Sompolinsky, PRX, 2021]. There, the analytical results for a deep linear network are heuristically extended to a network with ReLU activations, by replacing the linear GP kernels with the ReLU GP kernels. Despite having a more complex renormalization, also in our case the renormalized kernel can be seen as a linear combination, weighted by the order parameter, of many “GP” kernels, i.e., the path-path kernels. Therefore, we foresee how a similar argument could potentially be applied in our case. This suggestion, which was briefly mentioned in lines 330-332, will be discussed in more detail in the revised manuscript.\n\nAssumption (ii) is indeed a limitation of our model, when compared to the standard transformer. However, we would like to emphasize how this assumption allows us to consider an architecture featuring both multiple heads and multiple layers, enabling the characterization of attention paths interplay. The question is whether such a mechanism would still be present when relaxing assumption (ii), which seems reasonable to believe. The theoretical challenge, however, would be to disentangle the learning of attention paths interplay from the learning of the attention paths themselves, because now also the attention matrix would depend on the value weights. This discussion will be included in the revised manuscript. If the reviewer is interested, we provide further details in our response to reviewer 9bWs.\n\n2. *If I understand correctly, unlike the dense case, thanks to the transformer architecture, the kernel is not rescaled by a scalar renormalization variable, resulting in more interesting outcomes compared to [1]. Can you elaborate more on this?*\n\nThe reviewer is perfectly right: the multi-head transformer architecture induces a non-trivial renormalization encoded in the matrix order parameter. This results in a crucial improvement in performance. Indeed, for deep linear fully connected architectures, where the order parameter is purely scalar, the renormalization only affects the variance of the predictor, while the mean predictor remains the same as in the GP limit. Conversely, a matrix kernel renormalization produces a change in the mean predictor, thereby improving its generalization performance with respect to the GP limit. A matrix order parameter cleverly combines specific architectural components – the attention paths in our case  – offering deeper insight into the critical role of architecture in finite width networks. \n\nWe kindly invite the reviewer to also evaluate the importance of the above contribution for the theory of deep learning in general, independently of the transformer’s context. We elaborate on this point in our global response. \n\n3. *[...] The BPKR technique and its assumptions and rationale are not adequately introduced [...] the paper mentions the \"network’s kernel\" early on without really specifying which object it is considering [...].*\n\nWe thank the reviewer for this precious feedback, which helps us to improve the communication of our findings. The new version of the manuscript will integrate the reviewer’s suggestions by:\n- a. Specifying to which kernel we refer more precisely.\n- b. Outlining the rationale behind the BPKR method: explain how the expectation under the Gibbs distribution of network weights can be reduced to an expectation over a lower-dimensional distribution of macroscopic order parameters, which is obtained by gradual integration of the weights; explain how the so-reduced expectation can then be solved via the saddle-point method in the thermodynamic limit. \n- c. Clarify the significance of the thermodynamic limit assumptions.\n\n4. *It would be helpful to explain the structure of the different heads within the main text, rather than repeatedly directing the reader to long appendices.*\n\nWe welcome the reviewer’s suggestion to improve the clarity of Section 4.1.1. The revised manuscript will contain an explanation of the key features of the different heads, while keeping the mathematical definitions in the appendix. In particular, we will explain how the first good head makes use of the Markov nature of the task by attending only to nearby tokens and checking whether they match, while the second good head performs uniform attention.\n\n5. *In Fig. 3 (d), is the GD map obtained for the same network from which the trained query and key weights were obtained for the theory?*\n\nYes. The revised text will explain this more clearly.\n\n6. *In general, in your experiments, do you always apply attention to the input at all depths of the network?*\n\nThe attention matrix $\\Omega^{(\\ell)}$ is always a function of the bare input $x^{(0)}$, irregardless of the layer $\\ell$ (see Eq. 3). The attention matrix is however applied to the preactivation to that layer $x^{(\\ell)}$ (see Eq. 2).\n\n7. *What do you mean by “efficiently” (line 57)?*\n\nWe mean that we prune the less relevant heads so that the loss in performance is minimal. We have rewritten the sentence as follows: “we show that a trained network can be reduced in size with minimal performance loss.”\n\nWe hope our response provides clarifications to all the concerns the reviewer has raised. If you find our response convincing/useful, please consider increasing the score. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k5CvhrQRO6", "reviewer": "Reviewer_15FG", "summary": "This paper proposes a way to analyze the interplay of attention paths. Statistical analysis is provided, and experiments are performed to verify the theory.", "questions": "Please use formal theoretical claim to express the key quantify of interest.", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The idea considered in this paper is interesting, and the experiments seem to align with the theoretical justifications.", "weaknesses": "[1] This paper modifies the transformer model compared to the standard transformer. While changing the activation to linear seems to be acceptable for the ease of the derivation, the critical change in what is fed in the later layers seems strong: It significantly simplifies the format of equation (7) and the behavior of the considered transformer can be quite different from the standard transformer. The experiments are using two attention layers, which is significantly different from the real practice. It remains doubtful whether the analysis is still correct when using a deeper standard transformer.\n\n[2] The writing of this paper can be improved in its theory parts: There is no formal theory statement, and readers have to trust the authors that we use equation (9) - (13) in the analysis. Without the formal theory statement, it is even more difficulty to track the derivations in the appendix: Readers have to go through every detail to understand what the authors want to do in (9) - (13).\n\n[3] The writing of the paper can be improved in the experiment part: For example, it is not clear whether the authors implement a transformer using the considered architecture, or just a standard transformer. For image data, it is not clear whether additional embedding is used or not. In addition, the training details of the transformer is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.\n\nAt the moment, however, we believe that our contributions have not been evaluated in the proper context.\n\nIn particular, the reviewer seems to interpret our work as providing a method of analysis and its application. Instead, our work belongs to a more theory-focused line of research aimed at providing an analytical understanding of deep learning in the form of *exact analytical results*—in our case, an exact expression for the network's predictor on test examples.\n\nBefore reading this response, we kindly ask the reviewer to read our global response, which should help in better contextualizing our work.\n\n2. *[...] the critical change in what is fed in the later layers seems strong: It significantly simplifies the format of equation (7) and the behavior of the considered transformer can be quite different from the standard transformer. The experiments are using two attention layers, which is significantly different from the real practice. It remains doubtful whether the analysis is still correct when using a deeper standard transformer.*\n\nThe reviewer is correct in identifying the simplifications present in our model with respect to the standard transformer, which we openly acknowledge in our manuscript. However, we believe these simplifications should be considered in the context of our theory-focused line of research. At the current state of the art of deep learning theory, simplifications need to be made, if exact analytical results are to be derived. In this context, our work actually takes important steps forward with respect to the simplified models considered in other comparable works (see global response for details). In particular, although clearly a simplification, our assumption on what is fed to the attention heads allows us to analytically tackle a multi-head multi-layer architecture, which remained inaccessible to previous works. This allows us to characterize a previously uncovered learning mechanism: attention paths interplay.\n\nA very relevant question is whether this mechanism would still be present, when the attention is fed with the layer’s preactivation, rather than the bare input. It is reasonable to expect that the mechanism would persist, but that it would be harder to disentangle from the learning of the attention itself, which would now also depend on the value weights. Future work could build upon our results to tackle this exciting mathematical challenge. This consideration will be included in the revised version of the manuscript, as per request of reviewer Y9xT.\n\nRegarding the experiments, again the rigorous theoretical framework limits the scale of testable networks. In the Bayesian framework, one needs to sample from the high dimensional posterior distribution of network weights, which is much more computationally costly than gradient descent, and technically challenging for a larger number of layers/heads. Large scale experiments are anyways out of scope for this context. Since, under the specified assumptions, our analytical results are exact for any depth, experiments are typically of small scale and serve two purposes: 1) To convince the reader, who may not want to check the rigorous derivations, that the analytical predictions are exact 2) To provide minimal, easily interpretable examples illustrating the theoretical insights. We also invite the reviewer to check the scale of experiments considered in other comparable works, which we cite throughout the introduction.\n\n3. *There is no formal theory statement [...].*\n\nWe thank the reviewer for their suggestion to improve the readability of our results. The revised manuscript will provide a formal theory statement. Currently, our theoretical results and related definitions are presented alongside their interpretation and implications for generalization. The revised manuscript will first present the results (eqs. 9-12) and assumptions in a compact and more formal statement, and only later discuss their interpretation.\n\n4. *[...] it is not clear whether the authors implement a transformer using the considered architecture, or just a standard transformer. For image data, it is not clear whether additional embedding is used or not. In addition, the training details of the transformer is missing.*\n\nWe thank the reviewer for their important suggestions to improve the presentation of our experiments.\nWe always consider our transformer-like architecture, and never the standard transformer. This is because our work focuses on providing exact analytical results, rather than a method of analysis, so our experiments are focused on validating the theory and illustrating the predicted mechanisms on the considered architecture. The revised manuscript will make this clearer by explicitly referencing equations (1-4) defining the network, whenever it is mentioned in the experimental section. \n\nWhen nothing is said explicitly, the network is the one considered by the theory, i.e. “trained” by sampling its weights from the posterior distribution Eq. 8. The revised manuscript will state this explicitly. Details on the sampling method are given in “Experiments, Section F” in the Appendix. When the network is trained with gradient descent, we already state this explicitly and refer to  “Experiments, Section H.1” for training details.\n\nFor images, the input tokens consist of just the normalized pixel values. No additional embedding is used. This will be stated explicitly in the revised text.\n\nWe hope our response clarifies the scope of this work, and highlights our true contributions on both advancing the frontiers on the theoretical research on transformer-like models, and deep learning theory in general.  In our view, the current rating is unfair in light of these contributions.  If you find our response convincing/useful, please consider amending the score. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable time and their constructive comments. All reviewers agree on two main points:
- Our theoretical contribution is “interesting” (15FG and 5PZg), “novel” (Y9xT and 9bWs) and “technically sound” (Y9xT).
- Experiments thoroughly validate our theoretical results.

However, some reviewers raised concerns about the scale of the experiments (5PZg, 15FG) and the simplified architecture compared to the standard transformer (Y9xT, 15FG).  
While we acknowledge that our theoretical results involve several simplifications, we believe these are outweighed by the paper's novel contributions, especially when considered within the appropriate context and scope of this theory-focused work (as noted by 9bWs and Y9xT).

In this global response, we aim at better clarifying the context and scope of this work, and highlight important contributions currently overlooked in the reviews by 15FG and 5PZg.

We would first like to emphasize that our work is part of a line of research aimed at achieving a better analytical understanding of deep learning. It is important to underline that the focus is on obtaining *exact analytical results*—in our case, an exact expression for the network's predictions on test examples—rather than on applications. In particular, it is unfair to compare our results with those from heuristic and phenomenological approaches, which, even when inspired by theory, are less constrained by rigor and free to focus on any arbitrary architecture used in practice. 

Given the current state of the art in exact analytical methods, simplifying assumptions are typically necessary to obtain analytical results. In this context, we believe our work provides considerable contributions in the two following directions:
1. **Advances in the analytical understanding of learning in transformer-like architectures.** While from an application-focused perspective, our model still appears simplified with respect to the standard transformer, from the perspective of theory it integrates several key features that were not considered in previous works. These are the combination of (cf. Sec. 1 “Introduction” for references):
    - a. **A multi-head, multi-layer architecture.** Most previous works either considered single-head or single-layer models, which did not allow for a characterization of the interplay between attention paths.
    - b. **A large number of training examples.** Previous works considering the standard transformer architecture typically do so in the Gaussian Process (GP) limit, in which the network width $N→\infty$ while the number of examples $P$ stays finite (i.e. $P<<N$ in practice). Here instead we consider the more realistic case of $P=\alpha N$ for finite $\alpha$. As shown in our work, this is fundamental to uncover interesting learning mechanisms, which are lost in the GP limit.
    - c. **An exact derivation of the network predictor, independent of assumptions on the training data.** Even when including one of the features (a) or (b), previous works either provided results in the form of generalization bounds rather than exact expressions for the network predictor, and/or required specifying a model for the training data statistics. In contrast, our results apply to any realistic training dataset.

    While we openly acknowledge that deep learning theory has still a long way to go to analytically understand transformers, we believe to have made important steps forward in these still underdeveloped directions. Importantly, we believe that what we lose by making certain simplifications in our model is outweighed by what we gain in terms of the novel features described above. For example, note that the simplification of feeding the attention heads with the bare input (which is of concern for reviewers Y9xT and 15FG) is what allows us to analytically tackle a multi-head multi-layer architecture, thereby uncovering the mechanism of attention paths interplay. Therefore, while surely a limitation, this assumption allows us to make a significant step forward, on which future work could build upon, for example working on relaxing our assumption.
2. **Advances in the theory of deep learning in general.** While we focused on a transformer-like architecture due to its relevance to recent advancements in AI, our results have broader implications for the theory of deep learning. As mentioned above, exact derivations of the network predictor without assumptions on the dataset structure are typically limited to the network’s Gaussian Process (GP) limit. However, as shown in [Li, Sompolinsky, PRX 2021], it is only by going beyond this limit that the network’s hidden weights learn data dependencies. In the original deep linear network considered in [Li, Sompolinsky, PRX 2021], this phenomenon did not affect the mean predictor and only appeared as a scalar renormalization of the predictor’s variance. As correctly noted by Reviewer Y9xT, our work provides an analytically tractable yet reasonably powerful model in which the hidden weights learn a data-dependent structure that significantly enhances the performance of the mean predictor, which is now controlled by a matrix order parameter. There are currently only a few examples of such models in the literature [Li, Sompolinsky, NeurIPS 2022; Aiudi, et al., arXiv:2307.11807]. Furthermore our work provides a novel understanding of the order parameter as aimed at maximizing kernel-task alignment (the negative log-likelihood argument), allowing for a more direct interpretation of the renormalization phenomenon. In short, we believe our work makes an important contribution toward uncovering interesting learning mechanisms that can only emerge with a large number of training examples.

In synthesis, our work provides significant contributions to a theory-focused line of research that has always been of high relevance to NeurIPS, and which we hope to have better contextualized in our reply.

Meta-review: The paper studies thermodynamic limits of Transformers in a regime where the key-query matrices are fixed to some pre-trained forms. This leads to interesting new data-dependent kernels that exploit the structure of "attention paths" in order to show generalization benefits. The reviewers found this viewpoint interesting and relevant for understanding the benefits of Transformers.
Decision: Accept (poster)