Paper: JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention (ICLR-2024)

Reviewer comment (Evaluation):
Although the paper mentions experimental validation using real-world datasets and pre-trained models, it does not provide detailed quantitative evaluation metrics or comparisons against existing approaches. Including quantitative analysis would enhance the robustness and comprehensiveness of the findings.

Author response:
We thank the reviewer for the insightful and encouraging comments! 

Note that the reviewer may misunderstand our main contribution. We do not propose a novel approach to push SoTA, but instead we analyze existing Transformer architectures via a novel angle, by finding the joint dynamics of the self-attention and the MLP layer and using these new mathematical finding to analyze the training dynamics of multi-layer Transformer architectures. Therefore, there is no comparison against existing approaches. 

We indeed provide detailed quantitative analysis to show our theoretical findings are consistent with the real world scenarios. This includes the following:
+ We plot the change of entropy of attention patterns over time at different layers (Fig.6-8), as well as changes of stable rank in MLP layers (Fig. 7), verifying our theoretical findings (Sec. 4 and Theorem 4). 
+ We show that the hidden neurons in the Transformer MLP layer indeed learn the latent concepts in the hierarchy (Tbl. 1). 
+ We have multiple experiments that verify our intermediate theoretical findings. For example, Fig. 2 shows that Theorem 1 largely holds for softmax attention, even if the assumption for the theorem to hold for softmax may not be that realistic. Fig. 3 verifies Theorem 2 (growth in weights in the case of linear activations).

Reviewer comment (Reproducibility):
The paper does not delve into specific implementation details of the JoMA framework, such as hyperparameter choices or training procedures. Providing more information on these aspects would enable researchers to replicate the experiments and further validate the proposed approach.

Author response:
Since JoMA is not a novel approach but a way to analyze existing Transformer architectures with more realistic assumptions, we do not have specific implementation details of the framework like hyper-parameter choices. We follow the standard practice in training and mention our choices of hyper-parameters in Sec. 6.

Reviewer comment (Novelty):
Several lines of work around the learning preference of DNNs need to be discussed [cite1-3].

Author response:
# Comparison with existing works regarding learning preference of DNN

We thank the reviewer for providing these related work and will include the following comparison in our next revision. 

Both [cite1] and [cite2] focus on empirical study of the training dynamics of DNNs (not transformers) on real data vs random data, following the seminar paper “Understanding deep learning requires rethinking generalization” in ICLR’17. 

[cite1] shows that DNN will learn simple data points/patterns first, when training with real data, while for random data (either random input or random label), each data sample is equally difficult. It also studies the effects of different regularization on training curves on real/random data. [cite2] further shows that the concept of “simple samples/patterns” generalizes across multiple architectures and learning methods (e.g., SVM versus DNN). 

Note that both JoMA and [cite1][cite2] conclude that simple patterns are learned first, followed by learning of more complicated patterns. Both [cite1] and [cite2] treat DNNs as a black box and provide thorough empirical studies. They also do not analyze Transformers. In contrast, focusing on Transformer architecture that shows impressive performance across multiple domains, JoMA opens the architecture blackbox and gives more quantitative definition of patterns in terms of co-occurrence, and explains how the specific architecture of Transformer learns these patterns by implicitly learning a latent hierarchy. 

Furthermore, we could also relate our analysis to [cite1][cite2], by re-defining the somehow vague term “simple/difficult patterns” with more rigorous quantity: “simple patterns” can be co-occurred tokens in the lowest layer of hierarchy, while “difficult patterns” may come across multiple layers of hierarchy. In this case, our JoMA framework naturally explains why simple patterns are learned first, followed by difficult patterns. It also explains why the “easiness of pattern” translates from SVM to more complicated architectures: since co-occurred tokens in the lowest hierarchy layer can be easily learned by linear models. 

[cite3] focuses on empirical study of low (effective) rank bias of DNNs after training. It shows that low-rank bias happens in different optimizers and is insensitive to initializations. Compared to [cite1] and [cite2], [cite3] focuses less on the training dynamics but more on the property of the final trained model. In comparison, from the training dynamics point of view, JoMA not only provides a theoretical justification why such phenomena (i.e., emergence of low-rank structure) could happen during training, but also characterizes the change of ranks in a more refined way (i.e., high rank -> low rank -> high rank again) and discusses the underlying reason why the learning is done in this way (i.e., to learn salient feature first at lower layer, leaving non-salient features at the top layer of the hierarchy).

Reviewer comment (Novelty):
I am not sure how the insights from this paper could help us train/design better DNNs.

Author response:
# How the insights from JoMA help with practical applications for real-world scenarios
JoMA focuses on specific architectures and can provide meaningful suggestions on how the architectures can be changed. Here are a few examples:
+ Theorem 1 suggests that certain designs of soft-attention activation function lead to the invariance and future attention designs may be inspired by that. 
+ The attention entropy (and stable rank) re-bouncing curves suggest that low-rank structure may not be present over the entire training procedure. This suggests that we may want to use high-rank weight matrices at the beginning and the end of the training era. 
+ Fig. 8 shows that the re-bouncing of attention entropy over time is highly correlated with low validation error of the model after training. This suggests that we could check whether a run is healthy by checking the attention curve.  
+ From Sec. 5, we know that the learning of non-salient co-occurrence is slow, because the model is unsure whether they should be learned under the current hierarchy, or in the higher hierarchy. Knowing this, if we have prior knowledge on the structure of the feature hierarchy, we may develop novel algorithms to accelerate the training.

Reviewer **h5HX** also appreciates our work and thinks our understanding of hierarchical learning via self-attention can guide the design of follow-up transformer-based networks. 

Note that overall, the main focus of JoMA is to understand how multi-layer Transformer works and provide theoretical insights for the community, rather than improving practical applications for real-world scenarios.

Reviewer comment (Theory):
One assumption in the paper is that the embedding vectors of all tokens are orthogonal to each other. Does this need to hold for token embedding in every layer? Have the authors verified this assumption is valid in real cases? How much does this assumption affect the following conclusions?

Author response:
We thank the reviewer for the insightful feedback!

# Orthogonality assumption of embedding vectors
In order to reach concise conclusions, we made this assumption. The assumption needs to be true for all embedding layers (i.e., the upper out-projection matrix of MLP at every layer), otherwise there will be many additional terms in our analysis, and we won’t be able to give a clear overall picture of the training. 

In practice, we show that the embedding vectors are almost orthogonal to each other in many models. For Pythia models of different sizes (from 70M to 6.9B), the averaged absolute cosine similarity over all embedding vector pairs is bounded by 0.1 throughout the training, much smaller than the maximal value 1. For other pre-trained models like BERT, OPT-6.7B and ViT-Huge, it is also small (<0.05), for LLaMA-2-7B, it is <0.17. Please check Appendix B.1 for these additional experiments in our revised submission. 

Therefore, our analysis largely holds, but may contain many small terms. A detailed analysis with almost orthogonal vectors will be left for future work.

Note that for VIT, we measure the embedding layer (i.e., the upper out-projection matrix of MLP at every layer), but not the first input layer, which may have very strong correlations between nearby embeddings.

Reviewer comment (Experiments):
This paper analyze the training dynamics of multi-layer transformer and reveals how the attention is learned through time (first learning attention on tokens with most co-occurrence, and then expanding attention on other tokens). One question is, how is this dynamics dependent on model size? As we see from empirical results in the literature, larger model tends to learn attention that involves more complicated reasoning instead of simple co-occurrence. Can this factor be reflected by the analysis in the paper? How does the training dynamics (and entropy of attention) change when model size changes?

Author response:
# How larger models tend to learn attention that involves more complicated reasoning instead of simple co-occurrence
For this, we break the questions into two parts:
+ How do Transformers learn more complicated reasoning than simple co-occurrence?
+ How do large models help?

[**Learning more complicated reasoning**] In this work, we indeed focus on feature composition based on co-occurrence, which is a simple form (maybe the simplest form) of feature hierarchy. For reasoning, an intuition is that the model still captures co-occurring patterns. The difference is that thanks to embedding representation, these patterns are combinations of certain learnable “attributes” of input data, rather than input data themselves. 

For example, from the dataset “a x 3 = a a a”, “b x 3 = b b b”, the model will predict “c c c” from “c x 3 =”. Why can this be learned? From our point of view, there could be two critical components that contribute to this process:
+ The embedding of “a”, $u_a$, can contain multiple aspects of the letter “a”, e.g., $u_a = u_{[alphabet]} + u_{[first]}$, which means that a is “the first letter in the alphabet”. In this case, The pattern “[alphabet] x 3” is a common co-occurred pattern that can be extracted from “a x 3” and “b x 3”.  
+ As JoMA suggests, some neuron in the MLP hidden layer may represent the pattern p = “[alphabet] x 3”, and use an embedding to represent it for the next layer. Due to the residual connection, the resulting representation of “a x 3” can be $u_{p} + u_{[alphabet]} + u_{[first]}$. 
+ The co-occurred pattern $u_p + u_{[first]}$ will have a large inner product with the embedding “a” in the decoder layer, which also has a $u_{[first]}$ component. As a result, the letter “a” is decoded with high probability. Note that this procedure can generalize to any letter (e.g., c or d), as long as their embedding takes the form of $u_{[alphabet]} + …$.

In summary, for reasoning tasks:
+ Learning co-occurred patterns of the attributes that summarize different cases, by leveraging the power of embeddings. 
+ The matched pattern is also represented as embedding, superimposed to the original embeddings, thanks to residual connections. 
Overall, this gives very efficient representations and avoids enumerating all possible instantiation of patterns. Note that this is beyond JoMA's scope and we will leave a more mathematically rigorous study as the future work.

[**How do large models help?**] Large model helps in the following two ways:
+ Larger models typically have more hidden neurons in the MLP layers. According to lottery ticket hypothesis[1], it is more likely that there exists at least one hidden neuron that happens to initialize with a weight vector and an embedding that is useful for modeling a compositional concept. Such a good initialization will facilitate training. 
+ Larger models typically have larger feature dimension d, as suggested in Sec. 7 (Johnson–Lindenstrauss lemma), substantially more embedding vectors can be placed into a space with larger d, while maintaining almost orthogonality. 

Note that JoMA is just a starting point that captures the simplest mechanism of feature compositions. To make our analysis more straightforward and easier to follow, we also assume all embeddings are fixed/orthogonal during training. A thorough and more rigorous theoretical study of how Transformer performs reasoning, and how the embedding vectors are learned during training, is left for future work.

Reviewer comment (Theory):
Is the proposed JoMA also applicable to vision transformers? Since the input to vision transformer is different from text transformer in that the tokens are continuous instead of discrete and several assumptions in the paper may not hold. For example, the paper assumes tokens are orthogonal to each other, but for vision input lots of input tokens are very similar.

Author response:
We thank the reviewer for the insightful feedback!

# Orthogonality assumption of embedding vectors
In order to reach concise conclusions, we made this assumption. The assumption needs to be true for all embedding layers (i.e., the upper out-projection matrix of MLP at every layer), otherwise there will be many additional terms in our analysis, and we won’t be able to give a clear overall picture of the training. 

In practice, we show that the embedding vectors are almost orthogonal to each other in many models. For Pythia models of different sizes (from 70M to 6.9B), the averaged absolute cosine similarity over all embedding vector pairs is bounded by 0.1 throughout the training, much smaller than the maximal value 1. For other pre-trained models like BERT, OPT-6.7B and ViT-Huge, it is also small (<0.05), for LLaMA-2-7B, it is <0.17. Please check Appendix B.1 for these additional experiments in our revised submission. 

Therefore, our analysis largely holds, but may contain many small terms. A detailed analysis with almost orthogonal vectors will be left for future work.

Note that for VIT, we measure the embedding layer (i.e., the upper out-projection matrix of MLP at every layer), but not the first input layer, which may have very strong correlations between nearby embeddings.