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

We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.

At the moment, however, we believe that our contributions have not been evaluated in the proper context.

In particular, the reviewer seems to interpret our work as providing a method of analysis and its application. Instead, our work belongs to a more theory-focused line of research aimed at providing an analytical understanding of deep learning in the form of exact analytical results—in our case, an exact expression for the network's predictor on test examples.

Before reading this response, we kindly ask the reviewer to read our global response, which should help in better contextualizing our work.

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

The reviewer is correct in identifying the simplifications present in our model with respect to the standard transformer, which we openly acknowledge in our manuscript. However, we believe these simplifications should be considered in the context of our theory-focused line of research. At the current state of the art of deep learning theory, simplifications need to be made, if exact analytical results are to be derived. In this context, our work actually takes important steps forward with respect to the simplified models considered in other comparable works (see global response for details). In particular, although clearly a simplification, our assumption on what is fed to the attention heads allows us to analytically tackle a multi-head multi-layer architecture, which remained inaccessible to previous works. This allows us to characterize a previously uncovered learning mechanism: attention paths interplay.

A very relevant question is whether this mechanism would still be present, when the attention is fed with the layer’s preactivation, rather than the bare input. It is reasonable to expect that the mechanism would persist, but that it would be harder to disentangle from the learning of the attention itself, which would now also depend on the value weights. Future work could build upon our results to tackle this exciting mathematical challenge. This consideration will be included in the revised version of the manuscript, as per request of reviewer Y9xT.

Regarding the experiments, again the rigorous theoretical framework limits the scale of testable networks. In the Bayesian framework, one needs to sample from the high dimensional posterior distribution of network weights, which is much more computationally costly than gradient descent, and technically challenging for a larger number of layers/heads. Large scale experiments are anyways out of scope for this context. Since, under the specified assumptions, our analytical results are exact for any depth, experiments are typically of small scale and serve two purposes: 1) To convince the reader, who may not want to check the rigorous derivations, that the analytical predictions are exact 2) To provide minimal, easily interpretable examples illustrating the theoretical insights. We also invite the reviewer to check the scale of experiments considered in other comparable works, which we cite throughout the introduction.

3. There is no formal theory statement [...].

We thank the reviewer for their suggestion to improve the readability of our results. The revised manuscript will provide a formal theory statement. Currently, our theoretical results and related definitions are presented alongside their interpretation and implications for generalization. The revised manuscript will first present the results (eqs. 9-12) and assumptions in a compact and more formal statement, and only later discuss their interpretation.

4. [...] 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.

We thank the reviewer for their important suggestions to improve the presentation of our experiments. We always consider our transformer-like architecture, and never the standard transformer. This is because our work focuses on providing exact analytical results, rather than a method of analysis, so our experiments are focused on validating the theory and illustrating the predicted mechanisms on the considered architecture. The revised manuscript will make this clearer by explicitly referencing equations (1-4) defining the network, whenever it is mentioned in the experimental section.

When nothing is said explicitly, the network is the one considered by the theory, i.e. “trained” by sampling its weights from the posterior distribution Eq. 8. The revised manuscript will state this explicitly. Details on the sampling method are given in “Experiments, Section F” in the Appendix. When the network is trained with gradient descent, we already state this explicitly and refer to “Experiments, Section H.1” for training details.

For images, the input tokens consist of just the normalized pixel values. No additional embedding is used. This will be stated explicitly in the revised text.

We hope our response clarifies the scope of this work, and highlights our true contributions on both advancing the frontiers on the theoretical research on transformer-like models, and deep learning theory in general. In our view, the current rating is unfair in light of these contributions. If you find our response convincing/useful, please consider amending the score. Thank you very much.

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.

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

$

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

$

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

$

U=\underset{\tilde{U}}{\mathrm{argmin}} \ S(\tilde{U};X,Y), \qquad (11)

$

where the scalar function $S$ called the action is defined as

$

S(U;X,Y)=\mathcal{L}(U)+\alpha\mathcal{E}(U;X,Y), \qquad (12)

$

The scalar function $\mathcal{\mathcal{E}}$, which we call the energy, is given by

$

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

$

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

$

U ^{\pi\pi'}=\frac{1}{N}\mathbb{E}[V _{\text{eff}} ^{\pi}\cdot V _{\text{eff}} ^{\pi'\top}], \qquad (14)

$

where $V _{\text{eff}} ^{\pi}\in\mathbb{R} ^{1\times N}$ are the effective weights along path $\pi$ (Eq. 6).

Derivation: See Appendix C.

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

We thank the reviewer for their valuable time reviewing our work and for many positive comments. We’d like to address the reviewer’s remaining concerns as follows.

1. The analysis relies on several very strong assumptions: (i) linearity of the network output in the value weights, (ii) applying the attention at any depth on the network input, [...] the impact of (i) and (ii) on the conclusions remains unclear. It would be beneficial to discuss the implications of these assumptions further.

We agree with the reviewer that further discussion on assumptions (i) and (ii) would be beneficial.

Regarding assumption (i), we discuss in the paper a potential way forward, following the heuristic arguments in [Li, Sompolinsky, PRX, 2021]. There, the analytical results for a deep linear network are heuristically extended to a network with ReLU activations, by replacing the linear GP kernels with the ReLU GP kernels. Despite having a more complex renormalization, also in our case the renormalized kernel can be seen as a linear combination, weighted by the order parameter, of many “GP” kernels, i.e., the path-path kernels. Therefore, we foresee how a similar argument could potentially be applied in our case. This suggestion, which was briefly mentioned in lines 330-332, will be discussed in more detail in the revised manuscript.

Assumption (ii) is indeed a limitation of our model, when compared to the standard transformer. However, we would like to emphasize how this assumption allows us to consider an architecture featuring both multiple heads and multiple layers, enabling the characterization of attention paths interplay. The question is whether such a mechanism would still be present when relaxing assumption (ii), which seems reasonable to believe. The theoretical challenge, however, would be to disentangle the learning of attention paths interplay from the learning of the attention paths themselves, because now also the attention matrix would depend on the value weights. This discussion will be included in the revised manuscript. If the reviewer is interested, we provide further details in our response to reviewer 9bWs.

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

The reviewer is perfectly right: the multi-head transformer architecture induces a non-trivial renormalization encoded in the matrix order parameter. This results in a crucial improvement in performance. Indeed, for deep linear fully connected architectures, where the order parameter is purely scalar, the renormalization only affects the variance of the predictor, while the mean predictor remains the same as in the GP limit. Conversely, a matrix kernel renormalization produces a change in the mean predictor, thereby improving its generalization performance with respect to the GP limit. A matrix order parameter cleverly combines specific architectural components – the attention paths in our case – offering deeper insight into the critical role of architecture in finite width networks.

We kindly invite the reviewer to also evaluate the importance of the above contribution for the theory of deep learning in general, independently of the transformer’s context. We elaborate on this point in our global response.

3. [...] The BPKR technique and its assumptions and rationale are not adequately introduced [...] the paper mentions the "network’s kernel" early on without really specifying which object it is considering [...].

We thank the reviewer for this precious feedback, which helps us to improve the communication of our findings. The new version of the manuscript will integrate the reviewer’s suggestions by:

• a. Specifying to which kernel we refer more precisely.
• b. Outlining the rationale behind the BPKR method: explain how the expectation under the Gibbs distribution of network weights can be reduced to an expectation over a lower-dimensional distribution of macroscopic order parameters, which is obtained by gradual integration of the weights; explain how the so-reduced expectation can then be solved via the saddle-point method in the thermodynamic limit.
• c. Clarify the significance of the thermodynamic limit assumptions.

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

We welcome the reviewer’s suggestion to improve the clarity of Section 4.1.1. The revised manuscript will contain an explanation of the key features of the different heads, while keeping the mathematical definitions in the appendix. In particular, we will explain how the first good head makes use of the Markov nature of the task by attending only to nearby tokens and checking whether they match, while the second good head performs uniform attention.

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

Yes. The revised text will explain this more clearly.

6. In general, in your experiments, do you always apply attention to the input at all depths of the network?

The attention matrix $\Omega^{(\ell)}$ is always a function of the bare input $x^{(0)}$, irregardless of the layer $\ell$ (see Eq. 3). The attention matrix is however applied to the preactivation to that layer $x^{(\ell)}$ (see Eq. 2).

7. What do you mean by “efficiently” (line 57)?

We mean that we prune the less relevant heads so that the loss in performance is minimal. We have rewritten the sentence as follows: “we show that a trained network can be reduced in size with minimal performance loss.”

We hope our response provides clarifications to all the concerns the reviewer has raised. If you find our response convincing/useful, please consider increasing the score. Thank you very much.

We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.

At the moment, however, we believe that our contributions have not been evaluated in the proper context. Before reading this response, we kindly ask the reviewer to read our global response, which should help in better contextualizing our work.

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

As explained in the global response, our work is theory-focused and provides exact analytical expressions for the network predictor. In this context, experiments are typically designed to serve two purposes: 1) To convince the reader, who may not want to check the details of the theoretical derivations, that the analytical predictions are exact. 2) To provide minimal, easily interpretable examples illustrating the theoretical insights. Specifically, we focused on experiments with only 2 or 3 layers and 2 to 4 heads because the resulting order parameter remained visually interpretable in that case.

In the Bayesian framework adopted here, numerically validating our derivations on larger models and datasets would require an exponentially larger computational cost—much larger than with gradient descent—other than being technically challenging, because it requires sampling the network weights from an high dimensional posterior distribution (Eq. 8). We believe this to be out of the scope of our work: our theoretical results are rigorous and apply for any depth, under the thermodynamic limit specified by the theory.

We acknowledge that the term “deep” is used in the literature with varying meanings. In the context of theory, where analytical results have predominantly been obtained for single-layer attention, we followed the convention of using “deep” to emphasize that our derivations apply to models with an arbitrary number of layers. However, if the reviewer believes it is more appropriate, we are willing to change the term “deep” in the abstract to “multi-layer.”

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.

We thank the reviewer for spotting this problem. We have changed the order in both figures and made them more readable. The updated figures are available in the joint PDF. We have decided not to split the figures, because all the subplots are closely related and together show how the path combination mechanism affects the performance. For instance, in order to understand both the accuracy trend and the pruning mechanism in Fig. 3, it is useful to look at the structure of the order parameter.

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

We thank the reviewer for pointing us towards these missing references. These will appear in the revised version as follows. Refs. [Li et al., ICML 2024; Nichani et al., ICML 2024] will be included in the Introduction, in the discussion of other theory-focused, closely related works. Refs. [Olsson et al., 2022; Reddy et al. ICLR 2024] are also highly relevant, but more suitable for the Discussion section, since they focus on empirical methods [Olsson et al., 2022] and phenomenological models [Reddy et al. ICLR 2024] that are more distant from the focus of our work on exact derivations of the network’s predictor.

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

We thank the reviewer for pointing us towards this interesting work [1], which will be discussed in the Discussion section of the revised manuscript.

However, we would like to stress that paper [1] is very different in scope and methodology to our work. Our work is theory-focused, and aims at providing exact analytical expressions that describe the network’s learning mechanisms. The cited work [1] is application-focused, and aims at providing a competitive pruning algorithm. In our work, the application to head pruning is neither the main result nor the focus, and we are not making any claim about it being competitive: we only use it to elucidate the kernel combination mechanism and its interpretation.

We believe it unfair to compare our experiments to those appearing in [1]. As we explain in the global response, our work is part of a line of research aimed at achieving a better analytical understanding of deep learning. Given the current state of theory, simplifying assumptions are typically necessary to obtain exact analytical results. Given these constraints, it is not fair to compare the standard transformer used in practice to our simplified model. Real-world applications would anyways be out of scope for our work. We direct the reviewer to our global rebuttal, in which we explain how, when seen in the right context, our work presents novel theoretical insights, as well as several advancements compared to the simplifying assumptions of other closely related theoretical works.

We hope our response clarifies the scope of this work, and highlights our true contributions on both advancing the frontiers of the theoretical research on transformer-like models, and deep learning theory in general. In our view, the current rating is unfair in light of these contributions. If you find our response convincing/useful, please consider amending the score. Thank you very much.

[1] Liu et al., ICML 2023. Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

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

We thank the reviewer for their valuable time reviewing our work, for the many positive comments, and for the intriguing questions.

Below are our replies to the reviewer’s questions, as well as minor concerns.

1. One minor weakness is that the experiments are limited to binary classification.

We agree that considering a network with multiple outputs would be interesting, and readily implementable within our theory. We would like to share with the reviewer why we actively decided to not include these kinds of experiments in this first work exploring our theory. The renormalization effect coming from having a number M of multiple outputs is the same for any architecture (e.g. the same as in deep linear networks [Li, Sompolinsky, PRX 2021]), and results in the order parameter gaining an additional pair of “output indices” ranging from 1 to M. In this work, we wanted to focus on illustrating exclusively the renormalization effects that are proper of the transformer architecture—-that of attention paths combination. Exploring the interplay between the two renormalization effects (multiple outputs and multiple paths) could be the topic of future work, with potentially exciting applications to language and next-token prediction tasks. One caveat in this context, is that one would need to be careful of a too large number of outputs M, which may spoil the thermodynamic limit which assumes M finite.

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

It seems reasonable to believe that an attention paths combination mechanism would still persist in this case. The theoretical challenge, however, would be to disentangle the learning of attention paths interplay from the learning of the attention paths themselves. Indeed, when the attention is a function of the bare input, our results for the renormalized kernel $K=\sum_{\pi,\pi'\in\mathrm{\Pi}}U^{\pi\pi'}C_{\pi\pi'}$ show that the value weights’ statistics enter only in determining the order parameter $U$, that is the interplay between attention paths. If the attention were to be a function of the previous layer’s output, instead, also the attention path kernels $C$ would depend on the value weights, which in this case affect the learning of the attention paths themselves. In other words, the value weights would influence performance in two key ways: (i) by optimizing the representation of inputs to the attention heads, and (ii) by improving the recombination of attention paths.

Obtaining analytical progress in this scenario is challenging, primarily due to the introduction of non-linearities in the trained weights through the attention’s softmax. A first step to advance in this computation could be to consider linear attention heads, which may simplify the analysis.

This discussion will be included in the revised manuscript, as also prompted by a question from reviewer Y9xT.

3. In order to prune the attention paths, the order parameters 𝑈𝜋,𝜋′ (Eq. (13)) must be computed, which involves averaging over the Gibbs distribution (Eq. (8)). How can this be computed in real architectures?

We would like to be sure to consider the case meant by the reviewer with “real architectures”, so let us consider a few scenarios below:

• a. For our model in the Bayesian framework, the order parameter can be equivalently computed by either solving the self-consistent equation (as in figure 3d “theory”), or sampling the posterior (as in figure 3d “sampled”).
• b. For our model trained with gradient descent, one can just take the learned query and key weights, and plug them into procedure (a) to obtain the order parameter for the value weights. Another more empirical approach that we did not try, would be to get multiple “samples” of the value weights by retraining the network with gradient descent, but for fixed query and key weights coming from the first iteration of the training (this is somehow what we did to compute the order parameter in figure 3d “gradient descent”, only that there we used a single “sample”).
• c. For the more standard transformer with attention as a function of the bare input and trained with gradient descent, one could still do something similar to (b). The difference would be that now also the attention matrix $\Omega$ would depend on the value weights. Instead of taking the learned query and key weights, and plugging them into procedure (a), one could just take the entire attention matrix $\Omega$ (as computed on each example from the network's activations) and plug it into procedure (a).
• d. If we also add nonlinear MLP blocks in between each head, one could still apply the same procedure as in (c), but, when arriving to step (a), he could only proceed by sampling the posterior, while the option of solving the self-consistent equation would not be available due to the self-consistent equation not being defined for the nonlinear case. A possible way to define the self-consistent equation could be to heuristically extend it to the nonlinear case, using the same heuristic argument as in [Li, Sompolinsky, PRX, 2021]. There, the analytical results for a deep linear network are heuristically extended to a network with ReLU activations, by replacing the linear GP kernels with the ReLU GP kernels. Despite having a more complex renormalization, also in our case the renormalized kernel can be seen as a linear combination, weighted by the order parameter, of many “GP” kernels, i.e., the path-path kernels. Therefore, we foresee how a similar argument could potentially be applied in our case.

4. ...barring some steps like saddle-point methods

The revised manuscript will have an improved explanation of the saddle point technique. We ask the reviewer, if possible, to help us by further elaborating on the unclear points.

Once again, we thank the reviewer for this thorough review and the very positive rating.

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