The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training

We thank the reviewer for the thorough and positive review, as well as the valuable suggestions and questions.

First, we agree with the reviewer that we could have utilized the available manuscript space more effectively. In response to suggestions from Reviewers HwPP and JKhN, we will revise Section 2 by explicitly presenting Theorems 2.3 and 2.4, clearly stating their assumptions, conditions, and implications, integrating key points into the main text. We appreciate the reviewer’s feedback, as these adjustments will enhance both the rigor and clarity of our theoretical contributions.

Second, we address the questions raised by the reviewer:

1. The analysis presented in Section 2 is not specific to single-layer self-attention, but applies generally to any layer within a Transformer model. Empirically, however, we observe that early layers can deviate from our theoretical predictions compared to deeper layers. This is expected, as backpropagation causes early layers to receive noisier updates, leading to greater divergence from theoretical predictions. On the other hand, the connections between tokens $x_i$ and $x_j$ and the corresponding prediction errors encoded by $\beta_{ij}$ are accurately reflected in the weight updates of deeper layers (which are closer to the output). The degree of deviation depends on the model and dataset and is often minimal. Figure S4 (complementing Figure 3) shows examples of models where such deviations are negligible. Based on these observations, we hypothesize that symmetric initialization of early layers is unlikely to be beneficial. Indeed, we tested symmetric initialization of all layers of decoder-only models, and we did not observe any significant improvement in training speed compared to standard initialization. We refer to our answer to question 1 of reviewer jZNG for further details.

2. We thank the reviewer for pointing this out. Since Figure 2 presents the symmetry and directionality scores for pretrained models, we assume this comment refers to Figure 3. We will update the caption of Figure 3 to explicitly mention the model (Bert-Base-Uncased) and the number of layers (12).

We thank the reviewer for the careful, detailed, and valuable input. First, we address the weaknesses raised by the reviewer:

1. We acknowledge that we did not evaluate real-world deployments, as our goal was to analyze the structures that emerge in self-attention matrices during pretraining. However, our findings suggest potential downstream benefits worth exploring in future work. Indeed, symmetric initialization speeds up training and leads to lower final loss, which often correlates with better downstream performance.

2. We thank the reviewer for highlighting this point, and we agree that the clarity of some notations could be improved. We will revise our proofs accordingly, specifically Proposition A.10.

Next, we address the questions raised by the reviewer:

1. We have conducted symmetric initialization experiments on autoregressive models. Our observations indicate that these models quickly lose their initial symmetry, eventually converging to $W_{qk}$ matrices with low symmetry scores. Furthermore, we observed no significant improvement in training speed compared to standard initialization. For details, please refer to the plots of the loss curve and symmetry available here (12 layers model): https://drive.google.com/file/d/14Y8huSc7EajiLWiGjjQNM3-BPvRVCw-G/view?usp=sharing In contrast, as shown in the manuscript, encoder-only models clearly benefited from symmetric initialization, exhibiting faster training and higher overall symmetry scores.

2. We thank the reviewer for raising this interesting point. We have not yet explored how to extend our mathematical framework to cross-attention mechanisms. Nonetheless, our empirical results show that the encoder components of encoder-decoder models consistently exhibit higher symmetry scores than the decoder components, and this difference holds across models of varying sizes (see Figure S2). This empirical evidence suggests that Theorem 2.4 can potentially be extended to cross-attention.

3. Our results highlight that the implicit weight updates of the $W_{qk}$ matrix encode essential structures for Transformer training. We hypothesize that architectural designs that better preserve and reinforce these structures in both the column and row spaces of $W_{qk}$ could further enhance training efficiency.

On the other comments and suggestions:

1. We respectfully ask the reviewer for clarification regarding which additional layers should be visualized. Figures 3 and S4 already show symmetry and directionality scores across all layers of several models trained from scratch on three distinct datasets.

2. The main goal of our study was to characterize the structures that emerge in self-attention matrices during pretraining. As discussed above, future work should explore how symmetric initialization can be leveraged to potentially benefit from fine-tuning on downstream tasks.

We thank the reviewer for the detailed and constructive feedback provided.

Below, we address the reviewer’s questions, along with the related concerns highlighted as potential weaknesses.

1. Proposition 2.1 aimed to show that accurate token prediction depends on learning an effective bilinear form in the embedding space, represented by the implicit matrix product $\mathbf{W}_{qk}$. However, we agree with the reviewer that this proposition is not essential. We will revise Section 2 accordingly by: (a) removing Proposition 2.1 and integrating its key insights into the main text, and (b) expanding Section 2.3 and adding a new Section 2.4 to present Theorems 2.3 and 2.4 in greater detail, making better use of the available space.

2. We thank the reviewer for raising these points. The comments led to valuable additional experiments and analyses that helped strengthen the empirical support for our findings. To provide a clear response, we address the two questions separately:

• 2.1 We conducted additional experiments using larger models and observed a consistent training speed-up, in line with our previous results. Specifically, we trained a BERT-large model (24 layers, 3x more parameters than Bert-Base) on the same three datasets used for BERT-mini and BERT-base, that is, Jigsaw, Wikipedia, and Red Pajama. Although full training runs are set to 200k steps, the current results are based on intermediate checkpoints. For Wikipedia, after 123k steps, the loss decreased from 0.2151 (without symmetric initialization) to 0.1874 (with symmetric initialization), yielding a 56.5% speed-up. For Jigsaw, after 184k steps, the loss dropped from 0.8113 to 0.7612, with an 11.0% speed-up. Similarly, for Red Pajama, after 140k steps, the loss improved from 0.2261 to 0.2058, with a 37.8% speed-up. We will include the final results in Table 1 of the revised manuscript. Furthermore, we found a clear positive correlation between increased symmetry from symmetric initialization and both faster training and lower final loss. For details, please refer to: https://drive.google.com/file/d/1nuAGfyjVAp9suVL0NydVPOTQK56_AtvP/view?usp=sharing. Finally, while the reviewer correctly notes a smaller relative gain in speed-up from BERT-mini to BERT-base, this is expected due to neural scaling laws (e.g., Kaplan et al., 2020) that make improvements harder at larger scales.

• 2.2 We conducted preliminary experiments on training Vision Transformers on the CIFAR-10 and ImageNet-1k datasets. We did not observe a significant speed-up with these specific experiments. Nonetheless, we hypothesize that further analysis is necessary to check if symmetric initialization can speed up the training of vision Transformers.

3. We have conducted experiments to enforce symmetry across layers by adding a regularization loss term, as follows:

$L_{reg} = \frac{\| \mathbf{W}^l_{qk}\|}{\| {\mathbf{W}^l_{qk}}_s\|} \quad \forall l \in [0,L] ,$

where the denominator is the symmetric component of the $\mathbf{W}^l_{qk}$ matrix. However, this did not lead to noticeable improvements in training speed or final loss compared to the baseline. While we see training constraints that promote symmetry as a promising direction for future work, the specific regularization method we tested was not effective. For details on our results, please refer to: https://drive.google.com/file/d/1Paa7z6KxD11MdXyJyhU-xOUFqg6tJkSf/view?usp=share_link

4. We thank the reviewer for raising this important point. In our current work, we have successfully explored methods to exploit symmetry. We are investigating several approaches to leverage column dominance to improve autoregressive training. These ongoing experiments will be fully addressed and presented in a dedicated future study.

On the other comments and suggestions, we thank the reviewer for pointing out 3 typos in the manuscript. We will revise the manuscript accordingly.

We thank the reviewer for the careful review, for going through the proofs in the Appendix, and for the detailed feedback to improve clarity. While we appreciate the concerns raised, we respectfully disagree with the claim that the manuscript lacks the necessary mathematical rigor for its claims and propositions, and we outline our reasoning below.

First, we acknowledge the mistake in Equation (S8), but this has no impact on the correctness of our proofs or results. Indeed, like with Equation (S6), (S8) was intended as a direct factorization of the joint distribution under bidirectional training. However, it was included only to motivate Equations (S9) and (S10) (the standard Masked Language Modeling (MLM) objective) which is what we use throughout the paper, and does not depend on (S8). We will remove (S8) and keep only (S9) and (S10).

Second, we address the weaknesses regarding the explanation of our mathematical results (1-2). We value the reviewer’s feedback and have restructured the presentation of our results to improve clarity. We emphasize that our core results focus on deriving the implicit gradient of self-attention matrices, under minimal assumptions about the data. Importantly, the results in Sections 2.1–2.2 and Proposition A.7 rely solely on properties of self-attention. We are confident that the following revisions will better highlight the intuitions and clarify the theoretical setup:

• Section 2.1: We will keep the standard definition of self-attention and keep Equation (3) to connect it later with (8), (9), and Figure 1. Proposition 2.1 will be removed entirely, and only the essential content needed to link (3) to Proposition 2.2 will be preserved. Given these structural changes, we agree with the reviewer that Proposition 2.1 is not necessary for understanding the subsequent results. We will retain the reference to Section A.1 for relevant definitions and remove the proof of Proposition 2.1.

• Section 2.2: We will present a general definition of the negative log-likelihood in Equation (6) without introducing $C_i$ at that point. Then, in Proposition 2.2, we will define both $C_i$ and $P_j$, followed by a clear explanation of how the gradient of $W_{qk}$ can be derived using these two equivalent summations. This will provide the foundation for introducing the two main theorems.

• Section 2.3: We will enhance the explanation of how a token contributes to the gradient when used as context or prediction by explicitly including Proposition A.7 in this section. Up to this point, all mathematical results are derived purely from self-attention and do not assume anything about the input data. We will move the theorems to a new Section 2.4.

• Section 2.4: We will formally state Theorems 2.3 and 2.4 with their assumptions: (1) There are statistical correlations between tokens, a weak and general assumption that holds in most real-world data. As a result, token embeddings exhibit partial alignment, capturing semantic and predictive structure. Indeed, this alignment either exists in pretrained embeddings or naturally emerges during training in learned embeddings, encoding semantic relationships; (2) Entries of $W_{qk}$ are i.i.d. at initialization with finite mean and variance; (3) Bidirectional training induces approximately symmetric error signals, that is, the error in predicting token $i$ from $j$ is similar to that of predicting $j$ from $i$. This new section will clarify why the theoretical setup broadly applies to real-world data and what predictions it enables.

• Section 3-4: The current versions demonstrate how the predicted structures appear in Transformer models (Fig. 2 for language; Fig. S1 for vision and audio), emerge during training (Fig. 3), and can be leveraged in practical applications (Table 1). These results clearly show how our theoretical insights translate to real-world scenarios.

Finally, we address the questions raised by the reviewer:

1. Yes, during symmetric initialization, we randomly initialize $W_q$ and set $W_k = W_q$, ensuring $W_qW_k^T$ is symmetric, as we understand the reviewer suggests. We also experimented with a regularization term to enforce symmetry during training, but it did not improve training speed or final loss compared to the baseline. For details, see point 3 in our response to reviewer jKhN.

2. Our framework shows that, under bidirectional training, each token pair contributes to an approximately symmetric update of $W_{qk}$. However, it does not determine whether a fully symmetric $W_{qk}$ is "ideal," as this would require a precise definition of "ideal" and an analysis of attention matrices at convergence. This is an interesting direction for future work. Additionally, Remark A.15 highlights that, in MLM, only a subset of updates is symmetric. As a result, we naturally expect, and have empirically observed, non-fully symmetric $W_{qk}$, though these still exhibit significantly higher symmetry scores than those from autoregressive training.