JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention

Since its debut, Transformers (Vaswani et al., 2017) have been extensively used in many applications and demonstrates impressive performance (Dosovitskiy et al., 2020; OpenAI, 2023) compared to domain-specific models (e.g., CNN in computer vision, GNN in graph modeling, RNN/LSTM in language modeling, etc). In all these scenarios, the basic Transformer block, which consists of one self-attention plus two-layer nonlinear MLP, plays a critical role. A natural question arises:

How the basic Transformer block leads to effective learning?

Due to the complexity and nonlinearity of Transformer architectures, it remains a highly nontrivial open problem to find a unified mathematical framework that characterizes the learning mechanism of multi-layer transformers. Existing works mostly focus on 1-layer Transformer (Li et al., 2023a; Tarzanagh et al., 2023b) with fixed MLP (Tarzanagh et al., 2023a) layer, linear activation functions (Tian et al., 2023), and local gradient steps at initialization (Bietti et al., 2023; Oymak et al., 2023), etc.

In this paper, we propose a novel joint dynamics of self-attention plus MLP, based on Joint MLP/Attention Integral (JoMA), a first integral that combines the lower layer of the MLP and self-attention layers. Leveraging this joint dynamics, the self-attention is shown to have more fine-grained and delicate behavior: it first becomes sparse as in the linear case (Tian et al., 2023), only attends to tokens that frequently co-occur with the query, and then becomes denser and gradually includes tokens with less frequent co-occurrence, in the case of nonlinear activation. This shows a changing inductive bias in the Transformer training: first the model focuses on most salient features, then extends to less salient ones.

Another natural question arises: why such a learning pattern is preferred? While for 1-layer this does not give any benefits, in multilayer Transformer setting, we show qualitatively that such a dynamics plays an important role. To demonstrate that this is the case, we assume a hierarchical tree generative model for the input tokens. In this model, starting from the upper level latent variables (in which the top-most is the class label of the input sequence), abbreviated as $\texttt{LV}_{s}$, generates the latents $\texttt{LV}_{s-1}$ in the lower layer, until reaching the token level ($s=0$). With this model, we show that the tokens generated by the lowest latents $\texttt{LV}_{1}$ co-occur a lot and thus can be picked up first by the attention dynamics as “salient features”. This leads to learning of such token combinations in hidden MLP nodes, which triggers self-attention grouping at $s=1$, etc. In this way, the non-salient co-occurrences are naturally explained by the top level hierarchy, rather than incorrectly learned by the lower layer as spurious correlation, which is fortunately delayed by the attention mechanism. Our theoretical finding is consistent with both the pre-trained models such as OPT/Pythia and models trained from scratch using real-world dataset (Wikitext2 and Wikitext103).

We show that JoMA overcomes several main limitations from Scan&Snap (Tian et al., 2023). JoMA incorporates residual connections and MLP nonlinearity as key ingredients, analyzes joint training of MLP and self-attention layer, and qualitatively explains dynamics of multilayer Transformers. For linear activation, JoMA coincides with Scan&Snap, i.e., the attention becomes sparse during training.

Let the total vocabulary size be $M$, in which $M_{C}$ is the number of contextual tokens and $M_{Q}$ is the number of query tokens. Consider one layer in multilayer transformer (Fig. 1(b)):

Input/outputs. ${\bm{x}}=[x_{l}]\in\mathbb{R}^{M_{C}}$ is the input frequency vector for contextual token $1\leq l\leq M_{C}$, $1\leq q\leq M_{Q}$ is the query token index, $K$ is the number of nodes in the hidden MLP layer, whose outputs are $h_{k}$. All the quantities above vary across different sample index $i$ (i.e., $x_{l}=x_{l}[i]$, $q=q[i]$). In addition, $\phi$ is the nonlinearity (e.g., ReLU).

The Attention Mechanism. In this paper, we mainly study three kinds of attention:

• •

  Linear Attention (Von Oswald et al., 2022): $\sigma(x)=x$ and $A:=1$;

• •

  Exp Attention: $\sigma(x)=\exp(x)$ and $A:=\mathrm{const}$;

• •

  Softmax Attention (Vaswani et al., 2017): $\sigma(x)=\exp(x)$ and $A:={\bm{1}}^{\top}\left(\sigma({\bm{z}}_{q})\circ{\bm{x}}\right)$.

Here $\circ$ is the Hadamard (element-wise) product. ${\bm{b}}\in\mathbb{R}^{M_{C}}$ are the attention scores for contextual tokens, given by a point-wise attention function $\sigma$. $A$ is the normalization constant.

Embedding vectors. ${\bm{u}}_{l}$ is the embedding vector for token $l$. We assume that the embedding dimension $d$ is sufficiently large and thus ${\bm{u}}_{l}^{\top}{\bm{u}}_{l^{\prime}}=\mathbb{I}(l=l^{\prime})$, i.e., $\{{\bm{u}}_{l}\}$ are orthonormal bases. Let $U_{C}=[{\bm{u}}_{1},{\bm{u}}_{2},\ldots,{\bm{u}}_{M_{C}}]\in\mathbb{R}^{d\times M_{C}}$ be the matrix that encodes all embedding vectors of contextual tokens. Then $U_{C}^{\top}U_{C}=I$. Appendix B.1 verifies the orthogonality assumption in multiple pre-trained models (Pythia, LLaMA, etc).

Residual connections are introduced as an additional term ${\bm{u}}_{q}$ in Eqn. 1, which captures the critical component in Transformer architecture. Note that we do not model value matrix $W_{V}$ since it can be merged into the embedding vectors (e.g., by ${\bm{u}}_{l}^{\prime}=W_{V}{\bm{u}}_{l}$), while $W_{K}$ and $W_{Q}$ are already implicitly modeled by the self-attention logits $z_{ql}={\bm{u}}^{\top}_{q}W^{\top}_{Q}W_{K}{\bm{u}}_{l}$.

Gradient backpropagation in multilayers. In multilayer setting, the gradient gets backpropagated from top layer. Specifically, let $g_{h_{k}}[i]$ be the backpropagated gradient sent to node $k$ at sample $i$. For 1-layer Transformer with softmax loss directly applied to the hidden nodes of MLP, we have $g_{h_{k}}[i]\sim\mathbb{I}(y_{0}[i]=k)$, where $y_{0}[i]$ is the label to be predicted for sample $i$. For brevity, we often omit sample index $i$ if there is no ambiguity.

Note that this is true for layer-wise training: optimizing the weights for a specific Transformer layer, while fixing the weights of others and thus the statistics of backpropagated are stationary. For joint training, this condition also holds approximately since the weights change gradually during the training process. Under Assumption 1, Appendix A.1 gives an equivalent formulation in terms of per-hidden node loss.

Training Dynamics. Define the conditional expectation $\mathbb{E}_{q=m}\left[\cdot\right]:=\mathbb{E}\left[\cdot|q=m\right]$. Now let us consider the dynamics of ${\bm{w}}_{k}$ and ${\bm{z}}_{m}$, if we train the model with a batch of inputs that always end up with query $q[i]=m$, then:

Here $h^{\prime}_{k}:=\phi^{\prime}({\bm{w}}_{k}^{\top}{\bm{f}})$ is the derivative of current activation and ${\bm{g}}_{{\bm{f}}}:=\sum_{k}g_{h_{k}}h^{\prime}_{k}{\bm{w}}_{k}$.

Therefore, we don’t need to explicitly update self-attention, since it is already implicitly incorporated in the lower layer of MLP weight! For softmax attention, we verify that even with the assumption, the invariance proposed by Theorem 1 still predicts ${\bm{z}}_{m}(t)$ fairly well.

In nonlinear case, the dynamics turns out to be very different. In this case, $\Delta_{m}$ is no longer a constant, but will change. As a result, the dynamics also changes substantially.

Remarks. For linear attention, the ratio is different but the derivation is similar and simpler. Note that the convergence speed heavily depends on the magnitude of $\mu_{j}$. If $\mu_{j}>\mu_{k}$, then $\delta_{j}(t)\ll\delta_{k}(t)$ and $v_{j}(t)$ converges much faster than $v_{k}(t)$. Therefore, the salient (i.e., large) components is learned first, and the non-salient (i.e., small) component is learned later, due to the modulation of the extra term $\exp({\bm{v}}^{2}/2)$ thanks to self-attention, as demonstrated in Fig. 4.

A follow-up question arises: What is the intuition behind salient and non-salient components in $\bm{\mu}$? Note that $\bm{\mu}$ is an $\ell_{2}$-normalized version of the conditional token frequency ${\bm{x}}$, given the query $q=m$. In this case, similar to Theorem 2 (and Tian et al. (2023)), we again see that if a contextual token $l$ co-occurs a lot with the query $m$, then the corresponding component $\mu_{l}$ becomes larger and the growth speed of $v_{l}$ towards $\mu_{l}$ is much faster.

Relationship with rank of MLP lower layer. Since MLP and attention layer has joint dynamics (Theorem 1), this also suggests that in the MLP layer, the rank of lower layer matrix $W$ (which projects into the hidden nodes) will first drop since the weight components that correspond to high target value $\mu_{j}$ grow first, and then bounce back to higher rank when the components that correspond to low target value $\mu_{j}$ catch up later.

A critical difference between the training dynamics of linear and nonlinear MLP is that in the nonlinear case, although slowly, the non-salient components will still grow, and the entropy of the attention bounces back later. While for 1-layer Transformer, this may only slow the training with no clear benefits, the importance of such a behavior is manifested if we think about the dynamics of multiple Transformer layers trained on a data distribution generated in a hierarchical manner.

Consider a simple generative hierarchical binary latent tree model (HBLT) (Tian et al., 2020) (Fig. 7(a)) in which we have latent (unobservable) binary variables $y$ at layer $s$ that generate latents at layer $s-1$, until the observable tokens are generated at the lowest level ($s=0$). The topmost layer is the class label $y_{0}$, which can take $D$ discrete values. In HBLT, the generation process of $y_{\beta}$ at layer $s-1$ given $y_{\alpha}$ at layer $s$ can be characterized by their conditional probability $\mathbb{P}[y_{\beta}=1|y_{\alpha}=1]=\mathbb{P}[y_{\beta}=0|y_{\alpha}=0]=\frac{1}{2}(1+\rho)$. The uncertainty hyperparameter $\rho\in[-1,1]$ determines how much the top level latents can determine the values of the low level ones. Please check Appendix A.5 for its formal definition.

Remarks. If $y_{0}$ takes multiple values (many classes) and each class only trigger one specific latent binary variables, then most of the top layer latents are very sparsely triggered and thus $\rho_{0}$ is very close to $1$. If $\rho$ is also close to $1$, then for deep hierarchy and shallow common ancestor, $\mathbb{P}[y_{l}=1|y_{m}=1]\rightarrow 1$. To see this, assume $\rho=\rho_{0}=1-\epsilon$, then we have:

This means that two tokens $l$ and $m$ co-occur a lot, if they have a shallow CLA ($H$ small) that is close to both tokens. If their CLA is high in the hierarchy (e.g., $l^{\prime}$ and $m$), then the token $l^{\prime}$ and $m$ have much weaker co-occurrence and $\mathbb{P}(l^{\prime}|m)$ (and thus $x_{l^{\prime}}$ and $\mu_{l^{\prime}}$) is small.

With this generative model, we can analyze qualitatively the learning dynamics of JoMA: first it focuses on associating the tokens in the same lowest hierarchy as the query $m$ (and these tokens co-occur a lot with $m$), then gradually reaches out to other tokens $l^{\prime}$ that co-occur less with $m$, if they have not been picked up by other tokens (Fig. 5(b)); if $l^{\prime}$ co-occurs a lot with some other $m^{\prime}$, then $m$-$l$ and $m^{\prime}$-$l^{\prime}$ form their own lower hierarchy, respectively. This leads to learning of high-level features $y_{\beta}$ and $y_{\beta^{\prime}}$, which has high correlation are associated in the higher level. Therefore, the latent hierarchy is implicitly learned.

Dynamics of Attention Sparsity. Fig. 6 shows how attention sparsity changes over time when training from scratch. We use $10^{-4}$ learning rate and test our hypothesis on Wikitext2/Wikitext103 (Merity et al., 2016) (top/bottom row). Fig. 8 further shows that different learning rate leads to different attention sparsity patterns. With large learning rate, attention becomes extremely sparse as in (Tian et al., 2023). Interestingly, the attention patterns, which coincide with our theoretical analysis, yield the best validation score.

We also tested our hypothesis in OPT (Zhang et al., 2022) (OPT-2.7B) and Pythia (Biderman et al., 2023) (Pythia-70M/1.4B/6.9B) pre-trained models, both of which has public intermediate checkpoints. While the attention patterns show less salient drop-and-bounce patterns, the dynamics of stable ranks of the MLP lower layer (projection into hidden neurons) show much salient such structures for top layers, and dropping curves for bottom layers since they are suppressed by top-level learning (Sec. 5). Note that stable ranks only depend on the model parameters and thus may be more reliable than attention sparsity.

Validation of Alignment between latents and hidden nodes in MLP. Sec. 5 is based on an assumption that the hidden nodes in MLP layer will learn the latent variables. We verify this assumption in synthetic data sampled by HBLT, which generate latent variables in a top-down manner, until the final tokens are generated. The latent hierarchy has 2 hyperparameters: number of latents per layer ($N_{s}$) and number of children per latent ($N_{\mathrm{ch}}$). $C$ is the number of classes. Adam optimizer is used with learning rate $10^{-5}$. Vocabulary size $M=100$, sequence length $T=30$ and embedding dimension $d=1024$.

We use 3-layer generative model as well as 3-layer Transformer models. We indeed perceive high correlations between the latents and the hidden neurons between corresponding layers. Note that latents are known during input generation procedure but are not known to the transformer being trained. We take the maximal activation of each neuron across the sequence length, and compute normalized correlation between maximal activation of each neuron and latents, after centeralizing across the sample dimension. Tbl. 1 shows that indeed in the learned models, for each latent, there exists at least one hidden node in MLP that has high normalized correlation with it, in particular in the lowest layer. When the generative models becomes more complicated (i.e., both $N_{\mathrm{ch}}$ and $N_{l}$ become larger), the correlation goes down a bit.

Deal with almost orthogonal embeddings. In this paper, we focus on fixed orthonormal embeddings vectors. However, in real-world Transformer training, the assumption may not be valid, since often the embedding dimension $d$ is smaller than the number of vocabulary $M$ so the embedding vectors cannot be orthogonal to each other. In this setting, one reasonable assumption is that the embedding vectors are almost orthogonal. Thanks to Johnson–Lindenstrauss lemma, one interesting property of high-dimensional space is that for $M$ embedding vectors to achieve almost orthogonality $|{\bm{u}}_{l}^{\top}{\bm{u}}_{l^{\prime}}|\leq\epsilon$, only $d\geq 8\epsilon^{-2}\log M$ is needed. As a result, our JoMA framework (Theorem 1) will have additional $\epsilon$-related terms and we leave the detailed analysis as one of our future work.

Training embedding vectors. Another factor that is not considered in JoMA is that the embedding vectors are also trained simultaneously. This could further boost the efficiency of Transformer architecture, since concepts with similar semantics will learn similar embeddings. This essentially reduces the vocabulary size at each layer for learning to be more effective, and leads to better generalization. For example, in each hidden layer $4d$ hidden neurons are computed, which does not mean there are $4d$ independent intermediate “tokens”, because many of their embeddings are highly correlated.

Self-attention computed from embedding. JoMA arrives at the joint dynamics of MLP and attention by assuming that the pairwise attention score $Z$ is an independent parameters optimized under SGD dynamics. In practice, $Z=UW_{Q}W^{\top}_{K}U^{\top}$ is also parameterized by the embedding matrix, which allow generalization to tokens with similar embeddings, and may accelerate the training dynamics of $Z$. We leave it in the future works.

We propose JoMA, a framework that characterizes the joint training dynamics of nonlinear MLP and attention layer, by integrating out the self-attention logits. The resulting dynamics connects the dynamics of nonlinear MLP lower layer weights (projection into hidden neurons) and self-attention, and shows that the attention first becomes sparse (or weights becomes low rank) and then becomes dense (or weights becomes high rank). Furthermore, we qualitatively give a learning mechanism of multilayer Transformer that reveals how self-attentions at different layers interact with each other to learn the latent feature hierarchy.