Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers

Authors: lorenzo  tiberi
francesca  mignacco
kazuki  irie
haim  sompolinsky
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:32:07

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.

Main paper: 
 introduction: 
 In recent years, transformer models based on multi-head self-attention layers [1] [2] [3] [4] [5] [6] have achieved remarkable performance at natural language processing and vision tasks [7] [8] [9] . Yet, theoretical characterizations accounting for the success of these architectures remain sparse. Two fundamental questions remain to a large extent unsolved: First, interpretability-how can we discern task-relevant structures within the learned weights? Second, generalization-what specific aspects of the transformer architecture are responsible for their effective learning? We posit that one important feature of transformers is the combination of layer-wise multi-head organization with depth. This provides the network with a 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.In most cases, theoretical characterizations of transformers' expressivity [10] , inductive bias [11, 12] , generalization [13] [14] [15] and training dynamics [16] [17] [18] 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 in the first place, as in models consisting of a single-head [19, 15] , a single-layer [10, 11] , or both [12-14, 16, 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 [21] [22] [23] [24] , but address different questions than the present study, such as expressivity, generalization bounds, or phenomenological models. Further details on these and analogous works are discussed in Appendix I.One characterization of the complete transformer architecture has been obtained in the Bayesian framework under the infinite-width thermodynamic limit N → ∞ (and infinite number of heads H → ∞) [25, 26] , an actively studied regime in which neural networks become equivalent to Gaussian processes (GP) [27, 28] . However, the attention paths interplay is lost in this limit because the network's hidden weights remain statistically independent after learning. This limitation can be overcome by considering the finite-width thermodynamic limit [29] [30] [31] [32] , where also the number of examples P → ∞ such that P/N → α ∈ R + . In this regime, for example, multi-gated deep networks showcase task-relevant interplay between gates, mediated by the learned weights [33] .In this work, we apply the statistical mechanics theory of finite-width networks to a deep multihead 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 statistics 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 a task-relevant kernel combination mechanism, emerging beyond the GP limit (α > 0). This constructs the network's mean predictor as an 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 outside 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. 
model: 
 We consider a transformer-like [1] architecture consisting of a linear input projection layer; L multihead self-attention (MHA) layers, each having H attention heads; and a linear readout layer. The network input x ∈ R N0×T is a sequence of T tokens x t ∈ R N0 , with token index t ∈ {1, . . . , T }, and dimension N 0 . The input projection layer performs the transformationEQUATION)where N is the hidden layers' width. With the operator "•" we denote matrix-matrix or matrix-vector multiplication. The ℓ-th MHA layer with index ℓ ∈ {1, . . . , L} performs the transformationEQUATIONwhere, for each head h, we define the attention matrix Ω (ℓ)h ∈ R T ×T with matrix elements Task-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 is an improved kernel-task alignment in the finite-width regime (shown in (c)), enhancing generalization.EQUATION) = ( ) + ( ) + ( ) + ( ) + … K FW V (eff)π 1 ⊤ ⟩ ⟨V (eff)π 1 V (eff)π 2 ⊤ ⟩ ⟨V (eff)π 1 ⟨V (eff)π 2 ⟨V (eff)π 2 V (eff)π 2 ⊤ ⟩ V (eff)π 1 ⊤ ⟩ = ( 1 ) + ( 0 ) + ( 0 ) + ( 1 ) + … K GP μ = 1 μ = P μ = 2 … … … … V (eff)π 1 ⊤ V (eff)π 1 V (eff)π 2 ⊤ V (eff)π 2 V (eff)π 2 V (eff)π 1 ⊤ V (eff)π 1 V (eff)π 2 ⊤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 outputEQUATIONHere x t * can stand for different options for reducing the token dimension at readout, namely reading from a specific token t * or averaging over all tokens (x t * := 1 T x t ). The network's learnable parameters are the input projection weights V (0) ; the value, query and key weightsV (ℓ)h , W (ℓ)h Q , W (ℓ)h K L,H ℓ,h=1; and the readout weights a.Comparison to the Standard Transformer. The above architecture presents two main simplifications 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 as a direct function of the bare input x, rather than the processed input x (ℓ) . These simplifications allow us to apply back-propagating kernel renormalization (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 substantial: we will show that, in the finite-width regime, an important mechanism-task-relevant kernel combination-emerges, accounting for a 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. 1(a) ). We introduce the concept of attention paths, by defining a path "index" π := (h 1 , h 2 , . . . , h L ), where h 1 , . . . , h L ∈ {1, . . . , H}, which uniquely identifies each possible combination of the head indices across layers, i.e., each possible path through the attention heads. The network output can be rewritten asEQUATION)where Π is the set of all possible paths, and we define the "effective weights" asEQUATIONand the "attentioned input" asEQUATIONIn Eq. ( 5), the network can be seen as a deep linear network applied to a nonlinearly expanded input-the attentioned input ξ π . Through Eq. ( 7), we can see that the bare input x is nonlinearly expanded from an N 0 -dimensional space to an N 0 H L -dimensional space, by means of H L 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. 6). 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. 3, our theory characterizes the learned effective weights (Eq. 6) for a given, fixed realization of the query and key weights. 
theory: 
 A 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 [34, 35] . 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 of P example-label pairs {x µ , y µ } P µ=1 , we seek to characterize the Bayesian posterior [36] [37] [38] , or Gibbs distribution, over the parametersEQUATIONHere f (x µ , Θ) is the network output (Eq. 4) corresponding to the input x µ , where we emphasize its dependence on Θ, ∥•∥ is the Frobenius norm, σ 2 is the variance of the weights' Gaussian prior (set to σ = 1 throughout this paper), and T > 0 is the error variance, or Gibbs temperature (not to be confused with the number of tokens T ). Characterizing the Gibbs distribution allows to gain insights into the inductive bias imposed by the network architecture. Indeed, note that, in overparameterized networks, the Gibbs distribution for T → 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 configurations. For T > 0, 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. 2, we characterize the statistics of the weights Θ (the linear projection, value, and readout weights) for a fixed realization of the query and key weights. The fixed query and key weights can be given, for example, by pre-training the network with gradient descent, or by some task-informed initialization. In Sec. 4.2 we will show that the insights gained by our theory on the weights Θ 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. 3.1 below, we provide this formal result, accompanied by a sketch of its derivation and a discussion of the significance of the infinite-dimensional-or thermodynamic-limit under which our result is derived. In Sec. 3.2 we discuss the result's interpretation and its insights into the network's generalization capabilities.Definitions. Consider a training dataset consisting of P inputs x µ ∈ R N0×T and associated labels y µ ∈ R, where µ = 1, . . . P . Call X := {x µ } P µ=1 the set of training inputs and Y ∈ R P the vector of training labels with µ-th component y µ . Consider a network defined by Eqs. (1) (2) (3) (4) and in particular call f * the network output (Eq. 4) corresponding to a test input x * ∈ R N0×T .W (ℓ)h Q , W (ℓ)h K L,H ℓ,h=1are fixed, while all other weights Θ := V (0) , V (ℓ)h L,H ℓ,h=1 , a are distributed according to the Bayesian posterior distribu-tion defined in Eq. ( 8). Assume the "thermodynamic limit" N, N 0 , P → ∞, with P/N := α ∈ R + and P/(N 0 H L ) := α 0 ∈ R + , where α, α 0 as well as other size parameters T, H, L ∈ N are finite.Result 1. The mean predictor under the posterior distribution (Eq. 8) is given byEQUATIONThe vector k ∈ R P ×1 and the matrix K ∈ R P ×P , called training kernel, are defined in terms of a kernel function K : R N0×T × R N0×T → R as k µ := K (x * , x µ ) and K µν := K (x µ , x ν ), for µ, ν = 1, . . . , P . The kernel function is given byEQUATIONwhere ξ π (x) is the "attentioned input" corresponding to an input x ∈ R N0×T , along path π ∈ Π, as defined in Eq. (7) . The kernel function depends on a positive semi-definite matrix U ∈ R H L ×H L , called order parameter, which is given byEQUATIONThe scalar function S, called the action, consists of an "entropy" term L, and an "energy" termEQUATIONwhere K(U ; X) := K is the training kernel matrix. The expression for the entropy L is lengthy and is given in Appendix B.1. In the special case of H = 1, U is a scalar, andL (U ) = -σ -2(L+1) U + ln (U ).For general H, the entropy L (U ) is always maximized by U ππ ′ = σ 2(L+1) δ π,π ′ , which therefore is the solution of Eq. (11) in the GP limit defined by α → 0 + .Result 2. The matrix U obeys the following relationEQUATIONwhere V (eff)π ∈ R 1×N are the effective weights along path π , defined in Eq. ( 6).Derivation. See Appendix II.The derivation, which uses the BPKR technique [29, 33] , can be sketched as follows. Computing E [f * ] under the posterior distribution p (Θ) involves evaluating a high-dimensional integral in the weights Θ. The idea is to first reduce this computation into an integration over a lower-dimensional, 'macroscopic' variable U . Importantly, while Θ becomes infinite-dimensional as N → ∞, U remains finite-dimensional. The reduced integral is an expectation of the r.h.s. of Eq. ( 9), treated as a function of U , under the distribution p (U ) ∝ exp -1 2 N S (U ) , where S is the action defined in Eq. (11) . Then, this integral can be solved in the thermodynamic limit N → ∞, using the saddle-point method, which implies evaluating Eq. ( 9) at the U that minimizes the action (cf. Eq. 11). Crucially, the end result is fully characterized by this low-dimensional quantity U , commonly called order parameter in physics, which has a direct interpretation in terms of the network weights, given by Eq. (13) .In practice, the results obtained in the thermodynamic limit represent a good approximation also for the case of large but finite N . In this regard, the scaling of other hyperparameters with N → ∞ 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-width thermodynamic limit because in practice, for a given and typically large P , it is a good approximation only for very wide networks, when N ≫ P . In contrast, here we consider the finite-width limit in which P/N = α ∈ R + (which includes the GP limit for α → 0 + ). As can be seen from Eq. (11) , the action gains a new term for α > 0, which, as we shall discuss below, is fundamental to account for the learning of an attention paths interplay. Finally, we note that in our numerical experiments (Sec. 4) we will consider Bayesian networks which are overparameterized, i.e. P < N 0 H L , which is the network capacity at fixed query and key weights.Eq. ( 9) is a commonly found expression in thermodynamic theories of Bayesian learning, relating the network'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 kernel (Eq. 10) enables insights into how the transformer architecture favors this kernel-task alignment (Fig. 1(d) ). The kernel consists of the sum, weighted by the order parameter U , of many path-path kernels C ππ ′ , each computing the similarity between the attentioned input on two attention paths π and π ′ . 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 ∝ H L . 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 a total kernel with improved task alignment.This phenomenon, which we term task-relevant kernel combination, is indeed predicted by our results Eqs. (11) (12) . These state that the learned U minimizes a function S (Eq. 11), which, through the energy term E, favors kernel-task alignment. This can be seen by interpreting the energy term (Eq. 12) 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. 1(c) ).In contrast, in the GP limit α → 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 ππ ′ = σ 2(L+1) δ π,π ′ . Note that, in this limit, the benefits of kernel combination are lost (Fig. 1 (d), bottom line): First, out of all the path-path kernels C ππ ′ , only the same-path kernels (π = π ′ ) are used, while the cross-path kernels (π ̸ = π ′ ) 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] . A task-relevant kernel combination can therefore only emerge beyond the GP limit, in the finite-width regime α > 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 to make use of cross-path kernels, while controlling their magnitude allows to weigh the different path-path kernels in a task-relevant manner. 
experiments: 
 To corroborate our theoretical results, we "train" our model (Eqs. 1-4) by sampling its weights Θ (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 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 task-relevant kernel combination. In the second task, we will proceed to show analogous effects on classic image datasets (Omniglot [40] , MNIST [41] , 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).Task definition. The HMC classification task is defined as follows (Fig. 2(a) ). The µ-th example in the dataset corresponds to an hidden Markov chain q µ 1 , . . . , q µ T of length T = 30, alternating between two hidden states, q µ t ∈ {+, -}. The probability of transition to the opposite state (± → ∓) is p µ . The µ-th chain can belong to one of two classes, labeled y µ = ±1, depending on whether p µ = 0.3 or p µ = 0.7, respectively. The input tokens are a noisy, higher dimensional representation of the hidden states. These are given by x µ t = v q µ t + η µ t , where v ± ∈ R N0 are two orthogonal feature vectors corresponding to the states "±", with O (1) entries, while η µ t is a zero-mean Gaussian noise, with⟨η µ t η µ ′ ⊤ t ′ ⟩ = δ µ,µ ′ δ t,t ′ (σ 2 ∥ P ⊤ ∥ • P ∥ + σ 2 ⊥ P ⊤ ⊥ • P ⊥ ), where P ∥ and P ⊥ are the projectors along the subspace parallel or perpendicular to the plane spanned by v + and v -. Unless specified, σ ∥ = σ ⊥ = 1. The separate parameterization of the parallel (σ ∥ ) and perpendicular (σ ⊥ ) 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 on the classification accuracy, the second can always be filtered out by learning appropriate weights. We use P = 100 examples for training. We test the network performance in terms of the classification accuracy A = 1 P * µ δ y µ ,sign(⟨f µ ⟩) , where the sum is over a number P * = 1000 of test examples. Additional task details are given in Appendix G.1.We consider a network of L = 2 layers and H = 2 heads per layer, with readout from the first token. The network has a total of 4 attention paths, schematically depicted in Fig. 2(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 task-relevant kernel combination (for details, see Appendix G.2).We design the first head of each layer to give rise to a "good" attention path (green path) such that a network consisting of this good path alone achieves a high classification accuracy, A ∼ 94%. Along 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 to 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 is initialized randomly. This results in the three remaining paths having chance-level 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 to noisy data. In Fig. 2(c, top ) we show the network's classification accuracy as a function of the width N (blue, solid curve), compared to the GP limit (red, solid line). At lower N , well into the finite-width regime, we observe a considerable improvement in performance with respect to the GP limit. This can be understood in terms of an improved kernel-task alignment, as shown in Fig. 2(d) .x = + + + + + + + + ? + + + + t = 1 t = T c) d) b) e) f)v + v -v -v -v -v ? v ? v ? v ? v + v + v + a) "+1" "-1"This improved alignment is ensured by the order parameter U , plotted in Fig. 2(c, bottom ) for varying N . For N = 10, well into the finite-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 ππ ′ = δ π,π ′ , these paths are not suppressed, causing a deterioration in performance compared to that of the good path alone (red, dashed line in Fig. 2(c, top) ). The second benefit is apparent in the strong off-diagonals of U , anti-correlating the good and denoising paths. We 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 cross-path kernels. This additional improvement in performance becomes more apparent with noisier data. This is shown in Fig. 2(e) , where we plot the classification accuracy of the network consisting of only the good and denoising paths, on data with stronger perpendicular noise σ ⊥ = 5.Task definition. The one-shot image classification task (Fig. 3(a) ) is formulated in an in-context learning setting. The network is presented with a 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 i ∈ {1, . . . , p} of image a ∈ {1, 2, 3} corresponds to the token x (a-1)p+i , for a total of T = 3p tokens. We encode the labels +, -, ? using three fixed random vectors v + , v -, v ? ∈ R N0 , which we directly add to each patch (i.e., token) of the corresponding image. We also encode the token position with additive sinusoidal positional encoding [1] . The network is trained on the Omniglot dataset [40] , while we test its classification accuracy on both in-distribution (ID) unseen classes of Omniglot, and out-of-distribution (OOD) FashionMNIST dataset (we also report results on MNIST in Appendix H).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 L = 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, G = 128, and P = 528k (i.e., the entire training set from Omniglot). We refer to Appendix H.1 for further details on this process.The plots shown in Fig. 3 are analogous to those for the HMC task (Fig. 2 ), and illustrate analogous kernel combination phenomena. Fig. 3(b) shows the classification accuracy for varying N . Again, we observe a performance gap between the finite-width and GP regimes. Interestingly, this improvement in performance is preserved also OOD, on FashionMNIST. Again, Fig. 3(d) shows that the performance gap can be understood in terms of an 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. 3 (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 a stronger weighting of certain paths w.r.t. others. Interestingly, the block diagonal structure of the order parameter allows for a simple interpretation of the interplay between paths: correlations mostly occur between paths sharing the same head h 1 in the first layer, which also determines which paths are overall enhanced (h 1 = 1, 3) or suppressed (h 1 = 2, 4). This structure of the order parameter transfers qualitatively well also to the network trained with gradient descent. In Fig. 3 (c, "gradient descent") we show an empirical order parameter, obtained by computing Eq. 13 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 on 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 a 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 lowest-scoring heads are then pruned from the model. The head h at layer ℓ is assigned the scores (ℓ)h = π,π ′ ∈Π (ℓ)h |U π,π ′ |, where Π (ℓ)his the set of all paths passing through that head, and U is the order parameter derived from theory. Fig. 3(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 a marginal drop in performance. In Fig. 3 (f) 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. In Appendix H.2.3 we achieve an analogous size reduction of 25% on a larger model with H = 8 heads.Agreement with sampled statistics. Finally, we note that both figures 2(c,e) and 3(b,c) show good agreement of our theory with the mean predictor and order parameter sampled from Eq. 8. While our theory becomes exact in the N, P → ∞ limit, the agreement holds even for small N = 10. In particular, in figures 3(b,c), it holds even if N < H L H L + 1 /2 = 136, the number of independent entries in the order parameter, which is supposed to be a finite quantity in our theory. 
conclusion and discussion: 
 Conclusion. We introduce a transformer-like 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 demonstrate 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 a task-relevant kernel combination mechanism, where the network's total kernel results from the sum of many kernels, specific to pairs of paths, and optimally weighted to improve generalization. This mechanism is confirmed by experiments on both synthetic and real-world sequence classification tasks. More broadly, our results are relevant to the theory of deep learning, as they provide an example of non-scalar kernel renormalization [33, 43] in a widely adopted architecture such as the transformer, illustrating its importance in accounting for network performance. Non-scalar, as opposed to scalar [29] , renormalization can affect the network's mean predictor and therefore lead to improved generalization. Our work provides a novel interpretation of its benefits in terms of an optimized kernel-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 [45] . Our theoretical insights may therefore be relevant to the quest for minimalistic and resource-efficient models [46, 47] .Limitations and Outlook. The above results are enabled by our theory's ability to go beyond the GP limit [25] , as well as incorporating a multi-head, multi-layer architecture-a prerequisite for the very existence of attention paths. However, various limitations could still be addressed, opening for exciting research directions. For example, attention in 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. ( 10) 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 [48] [49] [50] . Note that, instead, our theory is readily extendable to the case of linear MLP blocks, as well as multiple outputs, following [29, 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.e., adding two new output indices to 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] [52] [53] .  
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 as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS 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. 
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 in this work.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and 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. 
 
 Appendix: 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 [11, 12] and the training dynamics [17, 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-to-sequence functions, while [21] highlights their computational limitations, demonstrating that self-attention cannot model periodic finite-state languages, nor hierarchical structure, unless the number of layers or heads increases with input length. [16, 22] derive error bounds for non-linear models, with trainable queries and keys. As in our work, [19, 10, 55] focus on the training of value matrices. In [19, 10] 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] [57] [58] [59] [60] . However, the works mentioned above derive generalization bounds and do not provide a tight characterization of the learning curves. A first tight analysis was done in [14] , considering a single-layer model with factored self-attention. The authors of [13] provide a tight analysis of a 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 are not defined. For instance, [44] finds that one transformer block can learn different linguistic tasks according to the position of the self-attention layer. [58] shows that a transformer block can learn to encode topical models. [55] shows that one transformer block can encode patch associations. A different line of works has studied attention layers in the infinite-width limit [25, 26] . In particular, [25] establishes an 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.We recall the definitions for our transformer model and theory, Sec. 2 and 3 of the main text.The hyperparameters are Gibbs temperature (theory only) σ :Gaussian prior variance (theory only)The network weights areV (0) ∈ R N ×N0 : input projection weights V (ℓ)h ∈ R N ×N : value weights a ∈ R 1×N : readout weights W (ℓ)h Q , W (ℓ)h K ∈ R G×N :query and key weights (fixed in the theory)where h = 1, . . . , H, and ℓ = 1, . . . , L.x ∈ R N0×T : input sequencex t ∈ R N0 : single token in the sequence where t = 1, . . . , T .We append an additional index µ = 1, . . . , P to quantities that refer to a specific example in the training set, such as, e.g.x µ ∈ R N0×T : µ-th example in the training set y µ ∈ R : µ-th training labelWe further defineX := {x µ } P µ=1 Y := {y µ } P µ=1the vector of training labelsThe network performs the following sequence of input transformations (Fig. 4 ). The input projection layer is defined as The output from each attention layer ℓ = 1, . . . , L is defined recursively as:EQUATIONEQUATIONwhere, for each head h, we define the attention matrix Ω (ℓ)h ∈ R T ×T with matrix elementsEQUATION)The function ζ denotes the softmax applied along the direction of the token indexed by s. The linear readout is defined asEQUATION)where we adopt the unified notationx (L+1) t * = 1 T T t=1 x (L+1) taverage over all the tokens x (L+1) t * readout from a single token at position t = t * (18) in order to consider different options of token readout, adopted in different tasks.We are interested in computing the posterior distribution over the network's weights Θ :=V (0) , V (ℓ)h L,H ℓ,h=1, a . This is analogous to the Gibbs distribution in statistical physics and is defined asEQUATIONwhere ∥•∥ is the Frobenius norm, and Z indicates the normalization, also called partition function.In the main text, we introduce the concept of attention paths, by defining a path indexEQUATIONwhich uniquely identifies each possible combination of the head indices, i.e. each possible path through the attention heads (Fig. 4 ). The network output can be written asf = 1 H L N N 0 π∈Π V (eff)π • V (0) • ξ π (21)where Π is the set of all possible paths, and we define the effective weights asEQUATIONand the attentioned inputξ π = T t0,...,t L-1 =1x t0 Ω(1)h1t0t1 Ω(2)h2 t1t2 . . .EQUATIONIn order to present our theoretical results and their derivation, it is useful to also introduce the following notation.Partial-path indices. We indicate withEQUATIONa collection of head indices, from layer ℓ up to layer L. For ℓ = 1, π 1 ≡ π is the path index defined above, Eq. 20, and in the main text. For ℓ ̸ = 1, π ℓ corresponds more generally to a partial path, starting from layer ℓ. We call the set of such paths Π ℓ . For ℓ = 1, Π 1 ≡ Π is the collection of all complete paths defined above and in the main text.Attentioned input at layer ℓ. We define the attentioned input at layer ℓ, for ℓ = 1, . . . , L, asEQUATIONwere we remind the reader of the compact notation t * defined in Eq. 18. The attentioned inputs at subsequent layers are related byEQUATIONRecall that ξ (ℓ-1)h ℓ-1 π ℓ := ξ (ℓ-1)h ℓ-1 h ℓ h ℓ+1 ...h L , according to the notation defined in Eq. 24.We also give specific definitions for the cases ℓ = 0 and ℓ = L + 1. For ℓ = 0, we define ξ (0)π1 ≡ ξ π , corresponding to the attentioned input defined above, Eq. 23, and in the main text. The following relation holdsEQUATION)For ℓ = L + 1 we defineEQUATIONHere we enunciate our theoretical results, for which a derivation is given in Appendix C and Appendix D.We derive (Appendix C) the statistics of the network prediction f (x * ) := f * on a new test example x * . These are given byEQUATIONwhere ⟨•⟩ denotes the expectation under the posterior distribution Eq. 19, and δf * = f * -⟨f * ⟩.The quantities K test ∈ R, k ∈ R P , and K ∈ R P ×P are defined in terms of a kernel function K : R N0×T × R N0×T → R. For µ, ν = 1, . . . , P we define K test := K (x * , x * ), k µ := K (x * , x µ ), and K µν := K (x µ , x ν ).We derive the following form for the kernelEQUATIONwhere ξ π (x) is the attentioned input corresponding to an input x, while U , called order parameter, is a path-by-path matrix of size |Π| × |Π|, where |Π| = H L denotes the size of the set Π. The value of the order parameter is determined self-consistently by minimizing a scalar function S, called action in the physics literature. This is defined asEQUATIONwhere, for any matrix M , we defineEQUATIONand E U (1) ; X, Y = 1 P ln det K(U (1) ; X)EQUATIONNote that K(U (1) ; X) is simply the kernel defined above, through Eq. 31, where we have explicitly written its dependence on U (1) and X. Here we have defined a collection of matrix order parameters U (ℓ) for ℓ = 1, . . . , L, of size |Π ℓ | × |Π ℓ |, were |Π ℓ | = H L+1-ℓ denotes the size of the set Π ℓ . We called U (1) ≡ U . The order parameter U (L+1) is instead a scalar. We also defined U U (ℓ+1)π ℓ ,π ′ ℓ ext = U (ℓ+1)h ℓ π ℓ+1 ,h ′ ℓ π ′ ℓ+1 ext := U (ℓ+1)π ℓ+1 ,π ′ ℓ+1 δ h ℓ ,h ′ ℓwhere π ℓ , π ′ ℓ ∈ Π ℓ and π ℓ+1 , π ′ ℓ+1 ∈ Π ℓ+1 are the partial-path indices defined in Eq. 24, while h ℓ , h ℓ ′ = 1, . . . , H.The action must be minimized w.r.t. all of the order parameters. Note that in the main text we defined S as a function of U (1) ≡ U alone. By this we mean S evaluated at its minimum for a fixed U (1) , so that only the minimization w.r.t. U (1) is left to perform.The self-consistent solution for the order parameters is obtained numerically by minimizing Eq. 32 with gradient descent methods. Details on the procedure are given in Appendix E.We derive (Appendix D) the following expression for the order parameter in terms of the network weightsEQUATION)where ⟨•⟩ denotes statistical averaging over the posterior distribution Eq. 19.Here we provide the derivation of our result for the predictor statistics, reported in Sec. B.1.Let us recall the posterior distribution (Eq. 19 )EQUATIONwhere Z indicates the normalization, also called partition function.It is useful to introduce a vector t ∈ R P of P auxiliary variables in order to linearize the squared error in Eq. 36. The partition function then readsEQUATION)where ı is the imaginary unit, t µ is the µ-th component of t, while f µ is shorthand for f (x µ , Θ), and we use the symbol "∝" to indicate that we are neglecting multiplicative constants. We also defined the Gaussian integration measures Dt = dt exp -T 2 ∥t∥ 2 and DΘ = dΘ exp -1 2σ 2 ∥Θ∥ 2 .Calling f * := f P +1 , we add a source term ∝ ıt P +1 f P +1 to the exponential in the partition function, corresponding to the test example x * := x P +1 , obtainingEQUATIONIn what follows we write Z ıt P +1 as Z, without writing its argument explicitly. The partition function Eq. 38 allows us to obtain the predictor statistics on the test example by differentiation ⟨f * ⟩ := f P +1 = d d (ıt P +1 ) ln (Z)EQUATION⟨δf * ⟩ := δf P +1 = d 2 d (ıt P +1 ) 2 ln (Z)EQUATIONwhere δf P +1 = f P +1 -f P +1 , and ⟨•⟩ denotes statistical averaging over the posterior distribution (Eq. 19) . In what follows, we will neglect any multiplicative constants in the partition function, that are independent of ıt P +1 , since they are irrelevant for the computation of Eq. 39 and Eq. 40.We proceed to compute the partition function (Eq. 38) by integrating all of the network weights Θ, and finally the auxiliary variable t. We use the back-propagating kernel renormalization (BPKR) method [29, 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 asEQUATION)where we defineEQUATIONBelow we focus on computing I(t). In what follows, we will write I (t) as I, without writing its argument explicitly.We start by integrating over the readout weights a. In Eq. 42, we substitute f µ = a • ξ (L+1)µ , obtainingEQUATIONIn what follows, it is useful to define q (L+1) = P +1 µ=1 t µ ξ (L+1)µ and in generalEQUATIONWe compute the Gaussian integral in the readout weights a, obtainingEQUATIONWith Θ (ℓ) we indicate the collection of weightsΘ (ℓ) := V (0) , V (ℓ ′ )h ℓ,H ℓ ′ ,h=1, i.e. all weights up to layer ℓ. 45 we haveL Plugging q (L+1) = 1 √ N H H h L =1 V (L)h L • q (L)h L into Eq.EQUATIONWe perform the integral over the value weights V (L)h H h=1 , by noticing that it is given by the product of N identical integrals of the formEQUATION)where we definedEQUATIONwith i, i ′ = 1, . . . , N and h L , h ′ L = 1, . . . , H. Here q(L)h L iindicates the i-th element of the vector q (L)h L ∈ R N . We may consider A as an N H × N H matrix, whose two indices run over the pairs (i, h L ) and (i ′ , h ′ L ). With this notation, the result of the value weights integration isEQUATIONUsing the matrix determinant lemma2 , we findEQUATION)where we defined the scalar R (L+1)EQUATION)We now enforce the identity of Eq. 51 by Fourier representation of the Dirac delta function, introducing the auxiliary scalar variable U (L+1) . We obtainEQUATIONEQUATIONIn the statistical physics language, R (L+1) and U (L+1) are called order parameters. Since N → ∞, we can solve the integral in R (L+1) with the saddle point method [61, 62] . The value of R (L+1) at the saddle point isEQUATIONTherefore, we obtainEQUATIONEQUATION)where we defined the "entropy" termEQUATIONan we introduced the|Π L | × |Π L | matrix U (L+1) ext, with matrix elementsEQUATION)Here π L , π ′ L ∈ Π L are the partial-path indices defined in Eq. 24, which in this case coincide with h L , h L ′ = 1, . . . , H. While the definition of U (L+1) ext may appear superfluous here, it is useful to perform the proof by induction in Sec. C.2.4.Next, we perform the integration over the weights at layer L -1. The steps are almost identical to 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.PluggingEQUATION)Once again we see that the integral in the value weights V (L-1)h H h=1 is given by the product of N identical integrals of the formEQUATIONwhere we definedEQUATIONwith i, i ′ = 1, . . . , N and h L , h ′ L = 1, . . . , H. Again, we may consider A as an N H × N H matrix, whose two indices run over the pairs (i, h L-1 ) and i ′ , h ′ L-1 . With this notation, the result of the value weights integration isEQUATIONUsing the matrix determinant lemma 3 , we findEQUATIONwhere we introduced the L) , with matrix elements defined as|Π L | × |Π L | matrix R (EQUATION)With the same procedure as in Sec. C.2.2, we introduce the order parameter R (L) and its conjugate U (L) , also a |Π L | × |Π L | matrix. We haveEQUATIONEQUATION)Again, we solve the integral in R (L) with the saddle point method. The value of R (L) at the saddle point isR (L) = σ 2 U (L)-1 -U (L+1)-1 ext(64) Therefore we obtainEQUATIONEQUATIONwhere we give a more general definition of the entropy Eq. 55, such that it can take a matrix argumentL U (L) •U (L+1)-1 ext = - 1 σ 2 Tr U (L) •U (L+1)-1 ext + ln det U (L) •U (L+1)-1 ext (66)and we introduced the|Π L-1 | × |Π L-1 | matrix U (L)ext , with matrix elementsEQUATION)Here π L-1 , π ′ L-1 ∈ Π L-1 and π L , π ′ L ∈ Π L and are the partial-path indices defined in Eq. 24, while h L-1 , h ′ L-1 = 1, . . . , H.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 formEQUATIONEQUATION)Here we defined a collection of matrix order parameters, one for each integrated layer. The order parameter U (ℓ) is a partial-path-by-partial-path matrix of size|Π ℓ | × |Π ℓ |, were |Π ℓ | = H L+1-ℓdenotes the size of the set Π ℓ . We also defined U EQUATION)where π ℓ , π ′ ℓ ∈ Π ℓ and π ℓ+1 , π ′ ℓ+1 ∈ Π ℓ+1 are the partial-path indices defined in Eq. 24, while h ℓ , h ℓ ′ = 1, . . . , H.Eq. 68 is verified for layer ℓ = L -1, which we derived in Sec. C.2.3. The induction step, integrating over the weights at layer ℓ -1, is done by plugging q (ℓ)π ℓ = H h ℓ-1 =1 V (ℓ-1)h ℓ-1 • q (ℓ-1)h ℓ-1 π ℓ into Eq. 68 and applying exactly the same steps presented in Sec. C.2.3.After integrating all of the network weights, we haveEQUATION)withEQUATIONEQUATIONplugging q (0)π1 = P +1 µ=1 ξ (0)π1,µ t µ into Eq. 71, we see that we need to perform the following integral in t ∈ R P exp -EQUATIONwhere for convenience we report here the kernel definitions given in Sec. B.1. The quantities K test ∈ R, k ∈ R P , and K ∈ R P ×P are defined in terms of a kernel function K : R N0×T × R N0×T → R. For µ, ν = 1, . . . , P we define K test := K (x * , x * ), k µ := K (x * , x µ ), and K µν := K (x µ , x ν ). The form of the kernel isEQUATIONwithEQUATIONComputing the Gaussian integral (Eq. 72) we obtainEQUATIONwhere we report here for convenience the definition of the action S given in Sec. B.1EQUATIONwithEQUATIONIn the limit N, P → ∞, P N → α ∈ R + , we solve the integrals in U (ℓ) L+1 ℓ=1 with the saddle point method [61, 62] . The partition function therefore takes the final formZ ∝ exp 1 2 ıt P +1 K test -k ⊤ • (K + T I) -1 • k ıt P +1 + ıt P +1 k ⊤ • (K + T I) -1 • Y(78) where we recall that K test , k, and K all depend on U (1) , which must be evaluated at the minimum of the action Eq. 76 with respect to all of its arguments U (ℓ) L+1 ℓ=1 . Differentiating by ıt P +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. B.1.Here we provide the derivation of our result on 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 as a continuation of Sec. C, to which we refer for definitions.For convenience, we report here the result we want to deriveEQUATIONwhere π 1 , π ′ 1 ∈ Π 1 are the path indices defined in Eq. 24, while ⟨•⟩ denotes statistical averaging over the posterior distribution Eq. 19. We also recall the definition of the network effective weightsEQUATIONAs in Sec. C.1, we start from the partition functionEQUATIONTo the partition function, we add a source term ∝ π1∈Π1 V (eff)π1 • q(1)π1 *, with qEQUATIONsuch that differentiating by q(1)π1 * allows us to obtainEQUATIONwhere q(0)π1 * ,iindicates the i-th component of the vector q(0)π1 * ∈ R N , while q (0) * ∈ R N ×H L is the matrix whose π 1 -th component is q (0)π1 * .As in Sec. C, we proceed to compute the partition function (Eq. 82) by integrating all of the network weights Θ, and finally the auxiliary variable t. We make the following observation. In Eq. 82, we can write explicitlyf µ = 1 √ N H L π1∈Π1 V (eff)π1 • ξ (1)π1,µ. Furthermore, as in Sec. C.2.1 we can define q (L+1) = P µ=1 t µ ξ (L+1)µ and in generalEQUATIONThen Eq. 82 takes the formEQUATIONRenaming q (1)π1 + q(1)π1 * → q (1)π1 , we see that the steps for computing Eq. 85 are identical to those in Sec. C, until the integration over the input projection weights V (0) .After integrating all of the network weights except the input projection V (0) we haveEQUATIONwithEQUATIONWe now substitute q (1)π1 → q (1)π1 + q(1)π1 * in Eq. 87, as well as q (1)π1 = 1EQUATIONEQUATIONThe integral in V (0) in Eq. 88 has the formEQUATIONwhere we definedEQUATIONEQUATION(1)π1 * ,i qEQUATIONwhere i = 1, . . . , N , while j, j ′ = 1, . . . , N 0 , and with the notation q(1)π1 * ,i and q(0)π1 i we indicate the i-th component of the vectors q(1)π1 * and q (0)π1 respectively.We may consider A j,j ′ as the elements of an N 0 × N 0 matrix A. With this notation, the partition function after performing the integral Eq. 89 isEQUATIONEQUATIONWe now differentiate the partition function in Eq. 92 by q(1) * , as specified by Eq. 83, obtainingEQUATIONThe remaining steps are the same as in Sec. C. We use the matrix determinant lemma 4 and the Woodbury matrix identity5 respectively to express ln det (A) and A -1 in terms of a |Π 1 | × |Π 1 | matrix R (1) with elementsR (1)π1,π ′ 1 = σ 2 N σ 2 N 0 H L q (0)π1⊤ • q (0)π ′ 1 π 1 , π ′ 1 ∈ Π Lwhose identity we enforce by Fourier representation of the Dirac delta function, introducing the auxiliary |Π 1 | × |Π 1 | matrix U (1) . The result of these operations isEQUATIONσ 2 N H L π1,π ′ 1 ∈Π1 U (1)π1,π ′ 1 q (0)π1⊤ • q (0)π ′ 1 (94)As in Sec. C, we solve the integral in R (1) with the saddle point method. The value of R (1) at the saddle point isEQUATIONPlugging this back into Eq. 94 we obtainEQUATIONThe calculation of the integral in t follows that in Sec. C.3. We obtainEQUATIONwhere ⟨•⟩ U (1) ∼p(U (1) ) denotes the expectation under the distribution p U (1) EQUATIONwhere the action S is the same defined in Sec. C.3, Eq. 76. Exactly as in Sec. C.3, we compute the expectation using the saddle point method, under the limit N, P → ∞, P N → α ∈ R + . We therefore arrive at the final resultEQUATION)where U (1) is the value at the saddle point of S.Part III Experiments E Numerical evaluation of the order parameterThe order parameter in our theory is defined self-consistently as the minimum of an action, Eq. 32. We determine this order parameter numerically, by minimizing the action using the Adam optimizer [63] . As we describe in Appendix H. We sample the network weights from the posterior distribution (Eq. 19) using Hamiltonian Monte Carlo sampling [64] . Specifically, we use the NumPyro implementation of the No U-Turn Sampler (NUTS) [65, 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 a total of 1000 samples. Due to limited computational resources, the number of samples is smaller and varies for N > 100, while we did not take samples at all for very large N ≥ 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 at a given N takes less than 12 hours on a single 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 T to not be relevant for improving the network's classification accuracy. We therefore set it to a small, but finite value of T = 0.01. In contrast, for the one-shot image classification task, we find tuning the temperature to 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.Here we give additional details on the hidden Markov chain (HMC) classification task.Here we recall the task definition, providing a few additional details. The µ-th example in the dataset corresponds to an hidden Markov chain q µ 1 , . . . , q µ T of length T = 30, alternating between two hidden states, q µ t ∈ {+, -}. The transition probability to the opposite state (" ± " → " ∓ ") is p µ . In other words, describing the "±" states as one hot vectors + : 1 0 and -: 0 1 , the transition probability matrix of the hidden Matkov chain is The µ-th chain can belong to one of two classes, labeled y µ = ±1, depending on whether p µ = 0.3 or p µ = 0.7 respectively.1 -p µ p µ p µ 1 -p µ .The network is presented with visible states -the input tokens -which are a noisy, higher dimensional representation of the hidden states. These are given by x µ t = v q µ t + η µ t Here v ± ∈ R N0 are two orthogonal feature vectors corresponding to the states "±". We set N 0 = 200 andv +,i = √ 2 if i ≤ 100 0 otherwise v +,i = 0 if i ≤ 100 √ 2 otherwiseThe term η µ t is a zero-mean Gaussian noise, with⟨η µ t η µ ′ T t ′ ⟩ = δ µ,µ ′ δ t,t ′ (σ 2 ∥ P ⊤ ∥ • P ∥ + σ 2 ⊥ P ⊤ ⊥ • P ⊥ ), where P ∥ and P ⊥ are the projectors along the subspace parallel or perpendicular to the plane spanned by v + and v -.We also preprocess the dataset in two ways: adding to each chain a beginning-of-sentence (bos) token of zeros at position t = 0, and concatenating to each token a one-hot positional encoding vector of size T + 1 (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 a small, but finite value of T = 0.01. A finite value of T is required to sample the predictor statistics, hence comparing the theoretical results with samples.Here we give the details of the initialization of the fixed query and key weights.We recall that we consider a network of L = 2 layers and H = 2 heads per layer, with readout from the first token (i.e., t * = 1). The network has a total of 4 attention paths (Fig. 5(a) ). For the first head of each layer, we make a good choice of the fixed query and key weights, which defines a "good" attention path, achieving a good classification accuracy (cf. Sec. 4.1 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 Ω (ℓ)h ∈ R T ×T for head h at layer ℓ. Its matrix elements are defined asΩ (ℓ)h st = ζ 1 N 0 √ G x ⊤ s • W (ℓ)h⊤ K • W (ℓ)h Q • x t , W (ℓ)h Q , W (ℓ)h K ∈ R G×(N0+T +1)where s, t = 0, . . . , T , while ζ is the softmax function, applied along the direction of the token index s, and G is the dimension of the query-key feature space. We directly initialize the query-key matrix productW (ℓ)h := W (ℓ)h⊤ K • W (ℓ)h Q . Note that W (ℓ)h is an [N 0 + (T + 1)] × [N 0 + (T + 1)] matrix, because we have appended a one-hot positional encoding vector to the N 0 -dimensional input tokens (see Sec. G.1). In order to define the heads, it is convenient to decompose W (ℓ)h into the block structure +1) acts only on the one-hot positional encoding subspace, W (ℓ)h ff ∈ R N0×N0 acts only on the subspace of the tokens' "features", andW (ℓ)h = β W (ℓ)h ff W (ℓ)h fp W (ℓ)h pf W (ℓ)h pp where W (ℓ)h pp ∈ R (T +1)×(TW (ℓ)h fp , W (ℓ)h⊤ pf ∈ R N0×(T +1) mix these two subspaces. The scalar 0 < β < ∞ is a parameter controlling the "hardness" of the softmax function (for β → ∞, the softmax becomes a hardmax). We set it to β = 10 for all heads.Let us define the good heads, i.e. W (ℓ)h for h = 1 and ℓ = 1, 2. 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, ℓ = 1 we defineEQUATION= W(1)1 pf = 0 andEQUATIONW (1)1 pp tt ′ is the component of W (1)1 pp at indices t,t ′ .The head h = 1, ℓ = 2 implements uniform attention Ω(2)1 st= 1 T +1 , ∀s, t = 0, . . . , T . It is defined by W (1)1 ff = W (1)1 fp = W (1)1 pf = 0, and W (1)1 pp = 1.As discussed in the main text, the attention path defined by the good heads achieves a good classification accuracy. We can give an intuition as to why this is the case. Intuitively speaking, in the limit of a hardmax attention (β → ∞) and no noise σ ∥ , σ ⊥ → 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 a 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-sentence 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.Let us define the random heads, i.e. W (ℓ)h for h = 2 and ℓ = 1, 2. These are initialized with Gaussian identically and independently distributed entriesW (ℓ)h pp i,j ∼ 1 N 0 N (0, 1) W (ℓ)h pp t,t ′ ∼ N (0, 1) W (ℓ)h pf t,j ∼ 1 √ N 0 N (0, 1) W (ℓ)h fp i,t ′ ∼ 1 √ N 0 N (0, 1)∀i, j = 1, . . . N 0 and ∀t, t ′ = 0, . . . , T . Note that we 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), 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 of why this is so by looking at their associated same-path kernels, Fig. 5 (b). We can see that both the good-path and the two adversarial-path kernels appear very structured, with sharp excursions in their values for different pairs of examples. However, while the good-path kernel structure appears to be aligned with the task, well distinguishing the two classes, the adversarial-path 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, as can be visually understood from the total GP kernel. In contrast, the total renormalized kernel, in which the adversarial-path kernels do not contribute, preserves the task-relevant structure. Differently, the denoising-path kernel appears less structured and more uniform, with weaker noisy excursions. In fact, what we 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 one implementing uniform attention and repeating the same experiment shown in Fig. 2 (f) in the main text. This is shown in Fig. 5 (c), were we plot the classification accuracy of the network consisting of the good and denoising paths alone, for the case of σ ∥ = 1 and σ ⊥ = 5. We can see that the results are completely analogous to those shown in Fig. 2 (f) in the main text.Here we 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 1028/172/432-splits for the train/validation/test class splits [67] as implemented in torchmeta [68] . We use the Adam optimizer [63] using an initial learning rate of 3e -4 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.Here we report further results on the one-shot image classification task.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 classification accuracy presented in Sec. 4.2 and below are therefore shown at the optimal temperature for the given N , obtained by scanning the set of temperatures T ∈ {a 10 -b , 1.0, 1.5}, where a ∈ {1, 2.5, 5, 7.5} and b ∈ {1, 2}.Note that temperature is optimized only for the in-distribution classification accuracy. In particular, we do not optimize for temperature when testing out-of-distribution, but rather keep the optimal temperature determined by evaluating the network in-distribution. For the network considered in the main text, we 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 as 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 * ⟩ = k ⊤ • (K + T I) -1 • Ywhere we recall thatK µν = 1 H L π,π ′ ∈Π U ππ ′ C µν ππ ′ µ, ν = 1, . . . , Pwhere C ππ ′ 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 a 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. 3 that this can be seen as the negative log-likelihood of the labels vector Y under a centered Gaussian distribution, whose covariance matrix is the kernel K. While the most effective way to minimize the energy term is that described in the main text (Sec. 3), 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). In all of our experiments, we always observe this phenomenon of growing magnitude to a certain extent. We perform the same experiments on the classification accuracy for a deeper network (L = 3, H = 3).The results are shown in Fig. 6 . Note that here we also report the test accuracy on MNIST, which was not shown in the main text. We 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, we can 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.We repeat the head pruning experiment for a 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, we see that the head scores determined by our theory are qualitatively in line with the performance loss caused by pruning the corresponding head (Fig. 7(c) ). Note that we do not expect a 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. 7 (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-distribution performance has only a marginal drop, identical to that obtained after pruning 2 heads in the smaller network considered in the main text (also accounting to 25% of its size, Fig. 7 (b)). We can also see that pruning up to three heads improves the out-of-ditribution performance on fashionMNIST, as we also observed for the smaller network.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are encouraged to reach out to the asset's creators.Our code is public: https://github.com/tiberilor/attention-paths-interplay 38th Conference on Neural Information Processing Systems (NeurIPS2024).det I + qq ⊤ = 1 + q ⊤ • q, where q is a vector.det I + Q • U • Q ⊤ = det I + U • Q ⊤ • Q , withQ and U being m×n and n×n matrices respectively.det I + Q • U • Q ⊤ = det I + U • Q ⊤ • Q , withQ and U being m×n and n×n matrices respectively.I + Q • U • Q ⊤ -1 = I -Q • I + U • Q ⊤ • Q -1 • U • Q T ,with Q and U being m × n and n × n matrices respectively. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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".  

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

3. Some references on the theoretical mechanism of Transformers in learning are missing.  

Olsson et al., 2022. In-context Learning and Induction Heads.

Li 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)

Nichani et al., ICML 2024. How Transformers Learn Causal Structure with Gradient Descent.

Reddy et al. ICLR 2024. The mechanistic basis of data dependence and abrupt learning in an in-context classification task.
questions: 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. 

Liu et al., ICML 2023. Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: (1) One minor weakness is that the experiments are limited to binary classification.
questions: 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?

(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: 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.
paper_summary: paper_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.

[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.
strength_weakness: strength_weakness: 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.
2. **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.
3. **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.
4. **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).
questions: 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?
6. 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?
7. In general, in your experiments, do you always apply attention to the input at all depths of the network?
8. What do you mean by “efficiently” (line 57)? 

[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: limitations: The paper adequately lists its limitations in Section 5. I do not foresee any potential negative societal impacts arising from this study.
paper_summary: paper_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.
strength_weakness: strength_weakness: [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.

[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).

[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.
questions: questions: Please use formal theoretical claim to express the key quantify of interest.
limitations: limitations: NA

Comments: title: Theory statement
Introductory comment to the theory statement
comment: Thank you for your reply. We would like to address your concerns and some misunderstandings.

As we explained in our previous comment, the model we consider is linear in the value weights, while it is nonlinear in the query-key weights (i.e. the attention paths). Since we apply the BPKR technique only to characterize the learning of the value weights, we are not making any approximation and our results are *exact* at any depth for the model under consideration. There is no “risk” of the theory breaking at larger depths since the network will always stay linear in the value weights, at any depth.

We also explained in our first rebuttal that simulating networks of larger depth is prohibitive in the Bayesian framework, but also out-of-scope for this theory-focused work since the theory is exact at any depth, and the experiments have the main purpose of illustrating its mechanisms in a minimal setting. We can however promise to attempt to simulate as-deep-as-possible networks for the appendix of the final version of the manuscript, though there is no reason to expect different results for the model under consideration. 

A related question raised by Reviewer Y9xT that we find relevant here is whether our theory could be extended to a network that is nonlinear in the value weights, which we do not consider in this work. This is an interesting mathematical challenge, which future work could tackle by building upon our theoretical advancements (as discussed in our Discussion, and in our dialogue with reviewer Y9xT). We kindly ask the reviewer to evaluate our contributions in terms of these theoretical advancements, which have been highlighted in our global response, as well as by reviewers Y9xT and 9bWs.  While we openly admit that deep learning theory is still far from analytically characterizing state-of-the-art transformers, we make several important steps forward compared to other related works.

Regarding the 3-layer network in Figure 6, we did not show the order parameter simply because the objective of the figure was different. We have the measurement of the order parameter for that network and, unsurprisingly given the exactness of the theory, we clearly see that for large alpha it contains strong off-diagonals, while for small alpha it is diagonal. This makes sense as the order parameter needs to account for the boost in performance for larger alpha which is shown in Figure 6. We will include these plots of the order parameter in figure 6.
Dear authors, 

Sorry for the late reply. I think using techniques from linear networks [23,26] to derive theoretical results in Transformers is risky when the network goes deep and becomes more nonlinear. Equations 11 and 12 explain that when $\alpha\rightarrow 0$, only the same-path kernels are used, while when $\alpha>0$, many cross-path kernels are used, which are verified in Figure 2(e). What if the model becomes deeper? The 3-layer experiments in Figure 6 do not include such results as Figure 2(e). If time allows, it is better to try 12-layer Transformers or 6-layer Transformers, since they are the size of GPT-2 base/small and are deep enough, which can help to address my concern.
Thank you very much for your response and the updated score. We are glad to hear that the reviewer found our response useful. We will make sure to include these discussions in the final version.
I thank the authors for their detailed answers to the reviews, which addressed my initial comments. I personally appreciate the theoretical physics approach of the work, particularly the formulation of the transformer kernel as a task-dependent combination of path-path kernels. Furthermore, I recognize that making (sometimes strong) assumptions is necessary to advance our theoretical understanding of deep learning systems. With the proposed changes, I recommend acceptance and have adjusted the score accordingly. I encourage the authors to include all the new discussions in the revised version.
We noticed a typo in our reply, point 3.c : 

"For the more standard transformer with attention as a function of *the bare input*" -> "For the more standard transformer with attention as a function of *its layer's preactivations*".
### Theory statement ###

**Definitions.** Consider a training dataset consisting of $P$
inputs $x ^{\mu}\in\mathbb{R} ^{N _0\times T}$ and associated labels
$y ^{\mu}\in\mathbb{R}$, where $\mu=1,\ldots P$. Call $X\coloneqq \textbraceleft x ^{\mu} \textbraceright _{\mu=1} ^{P}$
the set of training inputs and $Y\in\mathbb{R} ^{P}$ the vector of
training labels with $\mu$-th component $y ^{\mu}$. Consider a network
defined by Eqs. (1-4) and in particular call $f ^*$ the network
output (Eq. 4) corresponding to a test input $x ^* \in\mathbb{R} ^{N _0\times T}$.
We remind the reader of the following network hyperparameters: the
input's embedding dimension $N _0$, the hidden layer's width $N$,
the number of tokens $T$, the number of attention heads per layer
$H$, and the number of attention layers $L$.

**Assumptions.** Assume the network weights $\Theta\coloneqq\left(V ^{\left(0\right)},\textbraceleft V ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H},a\right)$
are distributed according to the Bayesian posterior distribution defined
in Eq. 8, with temperature $\mathcal{T}>0$, while the query and key
weights $\textbraceleft W _{Q} ^{\left(\ell\right)h},W _{K} ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H}$
are fixed. 

Assume $N,N _0,P\to\infty$, with $P/N\coloneqq\alpha\in\mathbb{R} ^{+}$
and $P/(N _0H ^{L})\coloneqq\alpha _0\in\mathbb{R} ^{+}$, where $\alpha$,
$\alpha _0$ as well as other size parameters $T,H,L\in\mathbb{\mathbb{N}}$
are finite.

**Results.** Under the above assumptions,

(1) the mean predictor under the posterior distribution (Eq. 8) is
given by
$$\mathbb{E}\left[f ^{*}\right]=k ^{\top}\cdot\left(K+\mathcal{T}\mathbb{I}\right) ^{-1}Y, \qquad (9) $$
where the average is w.r.t. to the posterior distribution (Eq. 8).

The vector $k\in\mathbb{R} ^{P\times1}$ and the matrix $K\in\mathbb{R} ^{P\times P}$
are defined in terms of a kernel function $\mathcal{K}:\mathbb{R} ^{N _0\times T}\times\mathbb{R} ^{N _0\times T}\to\mathbb{R}$
as $k ^{\mu}\coloneqq\mathcal{K}\left(x ^{*},x ^{\mu}|U\right)$ and
$K ^{\mu\nu}\coloneqq\mathcal{K}\left(x ^{\mu},x ^{\nu}|U\right)$, for
$\mu,\nu=1,\dots,P$. The kernel function is given by 
\begin{equation}
\mathcal{K}\left(x,x'|U\right)=\frac{1}{H ^{L}}\sum _{\pi,\pi'\in\Pi}U ^{\pi\pi'}C _{\pi\pi'}\qquad\mathrm{with}\qquad C _{\pi\pi'}\coloneqq\frac{1}{N _0}\xi ^{\pi}\left(x\right) ^{\top}\cdot\xi ^{\pi'}\left(x'\right)\, \qquad (10)
\end{equation}
where $\xi ^{\pi}\left(x\right)$ is the ``attentioned input'' corresponding
to an input $x\in\mathbb{R} ^{N _0\times T}$, along path $\pi\in\Pi$
(Eq. 7) and $\Pi$ is the set of all attention paths for a given architecture
with $|\Pi|=H ^{L}$.

The matrix $U\in\mathbb{R} ^{H ^{L}\times H ^{L}}$, called *order
parameter*, is a positive semi-definite matrix given by
\begin{equation}
U=\underset{\tilde{U}}{\mathrm{argmin}} \ S(\tilde{U};X,Y)\, \qquad (11)
\end{equation}
where the scalar function $S$ called the *action* is defined
as 
\begin{equation}
S(U;X,Y)=\mathcal{L}(U)+\alpha\mathcal{E}(U;X,Y)\, \qquad (12)
\end{equation}
The scalar function $\mathcal{\mathcal{E}}$, which we call the *energy*,
is given by 
\begin{equation}
\mathcal{E}(U;X,Y)=\frac{1}{P}\ln\det\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right)+\frac{1}{P}Y ^{\top}\cdot\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right) ^{-1}\cdot Y, \qquad (13)
\end{equation}
where $K\coloneqq K(X,X|U)$ is the $P\times P$ training kernel matrix,
defined according to Eq. (10). The expression for the scalar function
$\mathcal{L}$, which we call *entropy*, is lengthy and is given
in Appendix B.1. 

(2) In the particular case of a single head per layer $H=1$, $U$ is a scalar, and the entropy assumes the simple form $\mathcal{L}\left(U\right)=\sigma ^{-2\left(L+1\right)}U-\ln\left(U\right)$,
where $\sigma ^{2}$ is the variance of the Gaussian prior on the network
weights $\Theta$ (see Eq. 8). 

(3) For general $H$, $\mathcal{L}\left(U\right)$ is minimized by
$U ^{\pi\pi'}=\sigma ^{2\left(L+1\right)}\delta _{\pi,\pi'}$,
which therefore is always the solution of Eq. (11) in
the GP limit defined by $\alpha\to0 ^{+}$.

(4) The matrix $U$ obeys the following relation 
\begin{equation}
U ^{\pi\pi'}=\frac{1}{N}\mathbb{E}[V _{\text{eff}} ^{\pi}\cdot V _{\text{eff}} ^{\pi'\top}], \qquad (14)
\end{equation}
where $V _{\text{eff}} ^{\pi}\in\mathbb{R} ^{1\times N}$ are the effective weights along path $\pi$ (Eq. 6).

**Derivation:** See Appendix C.
3. *The presentation of Section 3 can be improved. It is better to present equations of Section 3 into Theorems/Propositions/Lemmas. Otherwise, it is difficult to find the point in this section.*

We thank you for your feedback, which helps us improve the communication of our findings. Currently the manuscript alternates the presentation of our results with their interpretation and implications for generalization performance. The revised manuscript will be reorganized to contain a paragraph with a formal theory statement, containing definitions, assumptions and results. The interpretation and discussion of such results will be given after such a statement. The statement is provided in the next comment.
Please refer to the submitted manuscript for any referenced equations which do not appear in the statement itself (i.e. Eqs 1-8, which will appear in the manuscript before the theory statement). Also note that we made the following changes in notation:
- 1. We renamed the query and key weights from $\textbraceleft Q ^{\left(\ell\right)h},K{} ^{\left(\ell\right)h}\textbraceright _{\ell,h=1} ^{L,H}$ to $\textbraceleft W _{Q} ^{\left(\ell\right)h},W _{K} ^{\left(\ell\right)h}\textbraceright _{\ell,h=1} ^{L,H}$, in order to avoid confusion with the kernel $K$.
- 2. We renamed the network predictor on a test example from $f ^{P+1}$ to $f ^{*}$.
- 3. We renamed the network inputs from $x ^{\left(0\right)}$ to $x$.

This new naming scheme will be consistently implemented throughout the revised manuscript.

As a minor note on terminology, please note that this work uses the approach of theoretical physics, which has a long-standing tradition in developing theories of neural networks. Following this approach, we provide our analytical results in the form of Results and Derivations, rather than Theorems and Proofs.
Thank you very much for taking the time to reassess the value of our work based on our theoretical results. We would like to address your specific questions below. For additional details, please also refer to our general response, where we provide an overview of our contributions, compared to previous comparable works.

2. *What are the theoretical challenges you overcome in Section 3? In other words, why existing methods cannot solve such a problem?*

Currently, works aiming at providing analytical results for learning in transformers need to consider simplified architectures and/or training regimes (cf. Introduction). The methods adopted to study such simplified models are typically well established (e.g. NNs as Gaussian Processes equivalence, teacher-student setting, dynamic mean-field theory, BPKR, etc...). The challenge, however, is to identify simplified models that are analytically tractable, yet rich enough to offer insights related to specific features of transformers.

In this landscape, the challenges we overcome are: 

- a. To consider a model beyond a single-head and/or single-layer architecture, allowing for the existence of attention paths—present in large number in actual transformers—and the characterization of their interplay.

- b. To characterize such a model beyond the Gaussian Process (GP) limit, in which the hidden weights remain random after learning, and therefore cannot learn any data-dependent structure, in particular any clever combination of attention paths.

By including the above features, our analysis reveals a previously uncovered mechanism of attention paths interplay, implemented by learned structures in the value weights, which are predicted and described in our theory by the order parameter. As noted by both reviewers 9bWs and Y9xT, these data-dependent structures remained inaccessible to previous studies which, even when considering multi-layer multi-head transformers, only did so in the GP limit.

1. *I roughly checked Appendix C. I do not think the analysis of Transformer architecture is special here. The references [23, 26] mentioned in line 131 are also for linear networks rather than Transformers. My concern is that the analysis here oversimplifies Transformer models into linear models, which weakens the theoretical novelty.*

The model we consider is not linear. It is linear in the value weights, but highly nonlinear in the query-key weights, through the attention operation. As we explain in the last part of Section 2, the model can be seen as a deep linear network in the value weights, which is however applied to a highly nonlinearly expanded input of dimension $N_0 H^L$: it reads from the $H^L$ attentioned inputs (i.e. the input transformed by a specific attention path), which are nonlinear functions of the original input of size $N_0$.

A more precise statement is that we exclusively characterize the learning of those weights in which the network is linear, i.e. the value weights (in this sense applying the BPKR technique found in [23, 26]). The other query-key weights, however, are still learned through gradient descent; simply, we do not attempt to describe their learning mechanism in our theory. 

The value weights learn to combine architectural features specific to the transformer—the nonlinear attention paths—showcasing a mechanism of attention paths combination that is specific to this work. Also note that the learning of the linear value weights is itself highly nonlinear: this is what gives rise to the nontrivial combination mechanism, which better aligns the network kernel with the task and improves generalization, as shown by our theory and experiments.

4. *What is the practical significance of this work in real-world applications?*

As a theory-focused research, we admit that the immediate implications of our results to practical applications are limited. Nevertheless, we'd like to believe that our work is a progress toward a more practically useful theory. Our order parameter is a simple measure of the learned weights, which provides useful information about the role and interplay of different attention paths. While work is still needed to extend its application to state-of-the-art networks, we illustrated its efficacy in our head pruning experiment (Sec 4.2.1) connecting our theoretical results in the Bayesian setting to more practical models trained with gradient descent.
Thank you for the response. It is acceptable that this work is theory-focused. Then, the contribution should be Section 3. I have the following questions since I hope the theoretical part to be novel. 
1. I roughly checked Appendix C. I do not think the analysis of Transformer architecture is special here. The references [23, 26] mentioned in line 131 are also for linear networks rather than Transformers. My concern is that the analysis here oversimplifies Transformer models into linear models, which weakens the theoretical novelty. 
2. What are the theoretical challenges you overcome in Section 3? In other words, why existing methods cannot solve such a problem? 
3. The presentation of Section 3 can be improved. It is better to present equations of Section 3 into Theorems/Propositions/Lemmas. Otherwise, it is difficult to find the point in this section. 
4. What is the practical significance of this work in real-world applications?
Thank you very much for taking the time to reassess the value of our contribution and for raising the score. While we hoped for a higher score, we respect your perspective and are willing to further improve our theoretical statement based on your feedback.

We understand that the reviewer may be referring to theoretical statements in the field of mathematics. In this case, we would like to clarify that our work adopts the methods and presentation style of theoretical physics, according to which our results have the validity of a predictive theory, rather than a mere analysis tool.

Theoretical physics has a long history in developing theories of artificial neural networks [1-5], whose contribution has always been of high relevance to conferences like NeurIPS (see, e.g., [6-11]). The field of theoretical machine learning is growing quickly, with different methods from mathematics, physics, and computer science, and each discipline has its own distinct goals and ways of communicating results.

We recognize the importance of making our work accessible to researchers across these disciplines and are committed to refining our presentation in this direction. We welcome any further feedback to help us improve our work.

### References ###
[1] Hopfield, J. J. Proc. Natl Acad. Sci. USA 79, 2554–2558 (1982)

[2] Amit, D. J., Gutfreund, H. & Sompolinsky, H. Phys. Rev. Lett. 55, 1530–1533 (1985)

[3] Gardner, E. J. Phys. A 21, 257–270 (1988)

[4] Gardner, E. & Derrida, B. J. Phys. A 21, 271–284 (1988).

[5] For a review on recent works, see, e.g.: Bahri, Y. et al. Ann. Rev. Cond. Matt. Phys. 11, 501–528 (2019)

[6] Krogh, Anders, and John Hertz. Advances in neural information processing systems 4 (1991).

[7] Cortes, Corinna, et al. Advances in neural information processing systems 6 (1993).

[8] Saxe, A., McClelland, J., & Ganguli, S. (2014). Proceedings of the International Conference on Learning Represenatations 2014.

[9] Bordelon, Blake, Abdulkadir Canatar, and Cengiz Pehlevan. International Conference on Machine Learning. PMLR, 2020.

[10] Gerace, Federica, et al.  International Conference on Machine Learning. PMLR, 2020.

[11] Lee, Jaehoon, et al. Advances in neural information processing systems 32 (2019).
I appreciate the authors in providing the detailed revised material for the theory. I still feel that it is not sth as we usually refer to as a "theoretical statement". However, if treating the results as a new analysis tool, it is an interesting contribution. I have raise my score to 5.
### Theory statement ###

**Definitions.** Consider a training dataset consisting of $P$
inputs $x ^{\mu}\in\mathbb{R} ^{N _0\times T}$ and associated labels
$y ^{\mu}\in\mathbb{R}$, where $\mu=1,\ldots P$. Call $X\coloneqq \textbraceleft x ^{\mu} \textbraceright _{\mu=1} ^{P}$
the set of training inputs and $Y\in\mathbb{R} ^{P}$ the vector of
training labels with $\mu$-th component $y ^{\mu}$. Consider a network
defined by Eqs. (1-4) and in particular call $f ^*$ the network
output (Eq. 4) corresponding to a test input $x ^* \in\mathbb{R} ^{N _0\times T}$.
We remind the reader of the following network hyperparameters: the
input's embedding dimension $N _0$, the hidden layer's width $N$,
the number of tokens $T$, the number of attention heads per layer
$H$, and the number of attention layers $L$.

**Assumptions.** Assume the network weights $\Theta\coloneqq\left(V ^{\left(0\right)},\textbraceleft V ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H},a\right)$
are distributed according to the Bayesian posterior distribution defined
in Eq. 8, with temperature $\mathcal{T}>0$, while the query and key
weights $\textbraceleft W _{Q} ^{\left(\ell\right)h},W _{K} ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H}$
are fixed. 

Assume $N,N _0,P\to\infty$, with $P/N\coloneqq\alpha\in\mathbb{R} ^{+}$
and $P/(N _0H ^{L})\coloneqq\alpha _0\in\mathbb{R} ^{+}$, where $\alpha$,
$\alpha _0$ as well as other size parameters $T,H,L\in\mathbb{\mathbb{N}}$
are finite.

**Claim.** Under the above assumptions,

(1) the mean predictor under the posterior distribution (Eq. 8) is
given by
$$\mathbb{E}\left[f ^{*}\right]=k ^{\top}\cdot\left(K+\mathcal{T}\mathbb{I}\right) ^{-1}Y, \qquad (9) $$
where the average is w.r.t. to the posterior distribution (Eq. 8).

The vector $k\in\mathbb{R} ^{P\times1}$ and the matrix $K\in\mathbb{R} ^{P\times P}$
are defined in terms of a kernel function $\mathcal{K}:\mathbb{R} ^{N _0\times T}\times\mathbb{R} ^{N _0\times T}\to\mathbb{R}$
as $k ^{\mu}\coloneqq\mathcal{K}\left(x ^{*},x ^{\mu}|U\right)$ and
$K ^{\mu\nu}\coloneqq\mathcal{K}\left(x ^{\mu},x ^{\nu}|U\right)$, for
$\mu,\nu=1,\dots,P$. The kernel function is given by 
\begin{equation}
\mathcal{K}\left(x,x'|U\right)=\frac{1}{H ^{L}}\sum _{\pi,\pi'\in\Pi}U ^{\pi\pi'}C _{\pi\pi'}\qquad\mathrm{with}\qquad C _{\pi\pi'}\coloneqq\frac{1}{N _0}\xi ^{\pi}\left(x\right) ^{\top}\cdot\xi ^{\pi'}\left(x'\right)\, \qquad (10)
\end{equation}
where $\xi ^{\pi}\left(x\right)$ is the ``attentioned input'' corresponding
to an input $x\in\mathbb{R} ^{N _0\times T}$, along path $\pi\in\Pi$
(Eq. 7) and $\Pi$ is the set of all attention paths for a given architecture
with $|\Pi|=H ^{L}$.

The matrix $U\in\mathbb{R} ^{H ^{L}\times H ^{L}}$, called *order
parameter*, is a positive semi-definite matrix given by
\begin{equation}
U=\underset{\tilde{U}}{\mathrm{argmin}} \ S(\tilde{U};X,Y)\, \qquad (11)
\end{equation}
where the scalar function $S$ called the *action* is defined
as 
\begin{equation}
S(U;X,Y)=\mathcal{L}(U)+\alpha\mathcal{E}(U;X,Y)\, \qquad (12)
\end{equation}
The scalar function $\mathcal{\mathcal{E}}$, which we call the *energy*,
is given by 
\begin{equation}
\mathcal{E}(U;X,Y)=\frac{1}{P}\ln\det\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right)+\frac{1}{P}Y ^{\top}\cdot\left(K(X,X|U)+\mathcal{T}\mathbb{I}\right) ^{-1}\cdot Y, \qquad (13)
\end{equation}
where $K\coloneqq K(X,X|U)$ is the $P\times P$ training kernel matrix,
defined according to Eq. (10). The expression for the scalar function
$\mathcal{L}$, which we call *entropy*, is lengthy and is given
in Appendix B.1. 

(2) In the particular case of a single head per layer $H=1$, $U$ is a scalar, and the entropy assumes the simple form $\mathcal{L}\left(U\right)=\sigma ^{-2\left(L+1\right)}U-\ln\left(U\right)$,
where $\sigma ^{2}$ is the variance of the Gaussian prior on the network
weights $\Theta$ (see Eq. 8). 

(3) For general $H$, $\mathcal{L}\left(U\right)$ is minimized by
$U ^{\pi\pi'}=\sigma ^{2\left(L+1\right)}\delta _{\pi,\pi'}$,
which therefore is always the solution of Eq. (11) in
the GP limit defined by $\alpha\to0 ^{+}$.

(4) The matrix $U$ obeys the following relation 
\begin{equation}
U ^{\pi\pi'}=\frac{1}{N}\mathbb{E}[V _{\text{eff}} ^{\pi}\cdot V _{\text{eff}} ^{\pi'\top}], \qquad (14)
\end{equation}
where $V _{\text{eff}} ^{\pi}\in\mathbb{R} ^{1\times N}$ are the effective weights along path $\pi$ (Eq. 6).

**Derivation:** See Appendix C.
We thank the reviewer for giving us the opportunity to provide the
formal theory statement, which we could not include in our original
response due to length constraints. The theory statement will appear in the next comment. As we stated in the rebuttal, all parts of the manuscript about the
interpretation and implications of our results will be put after this
theory statement. Please refer to the submitted manuscript
for any referenced equations which do not appear in the statement itself (i.e. Eqs 1-8, which will appear in the manuscript before the theory statement).
Also note that we made the following changes in notation:
- 1. We renamed
the query and key weights from $\textbraceleft Q ^{\left(\ell\right)h},K{} ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H}$
to $\textbraceleft W _{Q} ^{\left(\ell\right)h},W _{K} ^{\left(\ell\right)h}\textbraceright  _{\ell,h=1} ^{L,H}$,
in order to avoid confusion with the kernel $K$. 
- 2. We renamed the
network predictor on a test example from $f ^{P+1}$ to $f ^{*}$. 
- 3. We renamed the network inputs from $x ^{\left(0\right)}$ to $x$.

This new naming scheme will be consistently implemented throughout
the revised manuscript.
I appreciate the authors in preparing the rebuttal. Could you help provide a more detailed description for **"The revised manuscript will provide a formal theory statement."** It would be great to have the statement here.

Thanks.