SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuandong Tian<sup>1</sup> Yiping Wang<sup>2,4</sup> Beidi Chen<sup>1,3</sup> Simon Du<sup>2</sup>
<sup>1</sup>Meta AI (FAIR) <sup>2</sup>University of Washington <sup>3</sup>Carnegie Mellon University <sup>4</sup>Zhejiang University <sup>5</sup>yuandong,beidic}@meta.com, <sup>6</sup>ypwang61,ssdu}@cs.washington.edu
Transformer architectures have shown impressive performance in multiple research domains and have become the backbone of many neural network models. However, there is limited understanding on how Transformer works. In particular, with a simple predictive loss, how the representation emerges from the gradient training dynamics remains a mystery. In this paper, we analyze the SGD training dynamics for 1-layer transformer with one self-attention plus one decoder layer, for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a discriminative scanning algorithm: starting from uniform attention, it gradually attends more to key tokens that are distinct for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this scan and snap dynamics on synthetic and real-world data (WikiText).
The Transformer architecture [1] has demonstrated wide applications in multiple research domains, including natural language processing [2, 3, 4], computer vision [5, 6, 7], speech [8, 9], multimodality [10, 11], etc. Recently, large language models (LLMs) based on decoder-only Transformer architecture also demonstrate impressive performance [4, 12, 13], after fine-tuned with instruction data [14] or reward models [15]. Why a pre-trained model, often supervised by simple tasks such as predicting the next word [4, 3, 13] or filling in the blanks [2, 16, 17], can learn highly valuable representations for downstream tasks, remains a mystery.
To understand how Transformer works, many previous works exist. For example, it has been shown that Transformer is a universal approximator [18], can approximate Turing machines [19, 20], and can perform a diverse set of tasks, e.g., hierarchical parsing of context-free grammar [21], if its weights are set properly. However, it is unclear whether the weights designed to achieve specific tasks are at a critical point, or can be learned by SoTA optimizers (e.g., SGD, Adam [22], AdaFactor [23], AdamW [24]). In fact, many existing ML models, such as  $k$ -NN, Kernel SVM, or MLP, are also universal approximators, while their empirical performance is often way below Transformer.
To demystify such a behavior, it is important to understand the training dynamics of Transformer, i.e., how the learnable parameters change over time during training. In this paper, as a first step, we formally characterize the SGD training dynamics of 1-layer position-encoding-free Transformer for
next token prediction, a popular training paradigm used in GPT series [3, 4], in a mathematically rigorous manner. The 1-layer Transformer contains one softmax self-attention layer followed by one decoder layer which predicts the next token. Under the assumption that the sequence is long, and the decoder learns faster than the self-attention layer, we prove the following interesting dynamic behaviors of self-attention during training. Frequency Bias: it progressively pays more attention to key tokens that co-occur a lot with the query token, and loses attention to tokens that co-occur less. Discriminative Bias: it pays attention to distinct tokens that appear uniquely given the next token to be predicted, while loses interest to common tokens that appear across multiple next tokens. These two properties suggest that self-attention implicitly runs an algorithm of discriminative scanning, and has an inductive bias to favor unique key tokens that frequently co-occur with the query ones.
Furthermore, while self-attention layer tends to become more sparse during training, as suggested by Frequency Bias, we discover that it will not collapse to one-hot, due to a phase transition in the training dynamics. In the end, the learning does not converge to any stationary points with zero gradient, but ventures into a region where the attention changes slowly (i.e., logarithmically over time), and appears frozen and learned. We further show that the onset of the phase transition are controlled by the learning rates: large learning rate gives sparse attention patterns, and given fixed self-attention learning rate, large decoder learning rate leads to faster phase transition and denser attention patterns. Finally, the SGD dynamics we characterize in this work, named scan and snap, is verified in both synthetic and simple real-world experiments on WikiText [25].
Concurrent works on Transformer dynamics. Compared to [26] that uses  $\ell_2$  loss, our analysis focuses on cross-entropy, which is more realistic, imposes no prior knowledge on possible attention patterns inaccessible to training, and allows tokens to be shared across topics. Compared to [27] that analyzes "positional attention" that is independent of input data with symmetric initialization, our analysis focuses on attention on input data without symmetric assumptions. [28, 29, 30] give similar conclusions that self-attention attends to relevant tokens. In comparison, our work analyzes richer phenomena in 1-layer transformers related to frequency and discriminative bias, which has not been brought up by these works. For example, sparse attention patterns are connected with co-occurrence frequency of contextual token and query, characterization of such connection over training with softmax, including two-stage behaviors of attention logits, etc. We also leverage analytical solutions to certain nonlinear continuous dynamics systems that greatly simplifies the analysis. Detailed comparison can be found in Appendix B.
Expressiveness of Attention-based Models. A line of work studies the expressive power of attention-based models. One direction focuses on the universal approximation power [18, 31, 32, 33, 20]. More recent works present fine-grained characterizations of the expressive power for certain functions in different settings, sometimes with statistical analyses [34, 35, 36, 37, 21, 38, 39, 40]. In particular, there is growing interest in explaining the capability of in-context learning [41] of Transformer, by mapping the gradient descent steps of learning classification/regression into feedforward steps of Transformer layers [42, 43, 44, 45, 37, 46]. Different from our work, the results in these papers are existential and do not take training dynamics into consideration.
Training Dynamics of Neural Networks. Previous works analyze the training dynamics in multilayer linear neural networks [47, 48], in the student-teacher setting [49, 50, 51, 52, 53, 54, 55, 56, 57], and infinite-width limit [58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], including extensions to attention-based models [72, 73]. For self-supervised learning, works exist to analyze linear networks [74] and understand the role played by nonlinearity [75]. Focusing on attention-based models, Zhang et al. [76] study adaptive optimization methods in attention models. Jelassi et al. [27] propose an idealized setting and show the vision transformer [5] trained by gradient descent can learn spatial structure. Li et al. [26] show that the 1-layer Transformer can learn a constrained topic model, in which any word belongs to one topic, with  $\ell_2$  loss, BERT [2]-like architecture and additional assumptions on learned attention patterns. Snell et al. [77] study the dynamics of a single-head attention head to approximate the learning of a Seq2Seq architecture. While these papers also study the optimization dynamics of attention-based models, they focus on different settings and do not explain the phenomena presented in our paper.
Notation. Let  $\{\pmb{u}_k\}_{k=1}^M$  be  $d$ -dimensional embeddings, and  $\{x_t\}$  be discrete tokens. For each token,  $x_t$  takes discrete values from 1 to  $M$ , denoted as  $x_t \in [M]$ , and  $\pmb{x}_t \coloneqq \pmb{e}_{x_t} \in \mathbb{R}^M$  is the corresponding one-hot vector, i.e., the  $x_t$ -th entry of  $\pmb{x}_t$  is 1 while others are zero.  $\pmb{u}_{x_t}$  is the token embedding at location  $t$  in a sequence.
Let  $U = [\pmb{u}_1, \dots, \pmb{u}_M]^\top \in \mathbb{R}^{M \times d}$  be the embedding matrix, in which the  $k$ -th row of  $U$  is the embedding vector of token  $k$ .  $X = [\pmb{x}_1, \dots, \pmb{x}_{T-1}]^\top \in \mathbb{R}^{(T-1) \times M}$  is the data matrix encoding the sequence of length  $T-1$ .  $XU \in \mathbb{R}^{(T-1) \times d}$  is the sequence of embeddings for a given sequence  $\{x_1, \dots, x_{T-1}\}$ . It is clear that  $X\mathbf{1}_M = \mathbf{1}_{T-1}$ .
We use  $X[i]$  to denote  $i$ -th sample in the sequence dataset. Similarly,  $x_{t}[i]$  is the token located at  $t$  in  $i$ -th sample. Let  $\mathcal{D}$  be the dataset used for training.
1-Layer Transformer Architecture. Given a sequence  $\{x_{1},\ldots ,x_{T},x_{T + 1}\}$ , the embedding after 1-layer self attention is:
$$
\hat {\boldsymbol {u}} _ {T} = \sum_ {t = 1} ^ {T - 1} b _ {t T} \boldsymbol {u} _ {x _ {t}}, \quad b _ {t T} := \frac {\exp \left(\boldsymbol {u} _ {x _ {T}} ^ {\top} W _ {Q} W _ {K} ^ {\top} \boldsymbol {u} _ {x _ {t}} / \sqrt {d}\right)}{\sum_ {t = 1} ^ {T - 1} \exp \left(\boldsymbol {u} _ {x _ {T}} ^ {\top} W _ {Q} W _ {K} ^ {\top} \boldsymbol {u} _ {x _ {t}} / \sqrt {d}\right)} \tag {1}
$$
Here  $b_{tT}$  is the normalized self-attention weights  $(\sum_{t=1}^{T-1} b_{tT} = 1)$ . One important detail is that we mask the weight that the query token attends to itself, which is also being used in previous works (e.g., QK-shared architecture [78]). See Sec. 7 for discussions about residual connection. Let  $\boldsymbol{b}_T \coloneqq [b_{1T}, \dots, b_{T-1,T}]^\top \in \mathbb{R}^{T-1}$  be an attention vector, then  $\boldsymbol{b}_T^\top \mathbf{1} = 1$  and  $\hat{\boldsymbol{u}}_T = U^\top X^\top \boldsymbol{b}_T$ .
$\ell_2$ -Normalization. We consider adding a normalization to the output of the self-attention layer:  $\tilde{\pmb{u}}_T = U^\top \mathrm{LN}(X^\top \pmb{b}_T)$ , where  $\mathrm{LN}(\pmb{x}) \coloneqq \pmb{x} / \| \pmb{x}\|_2$ . NormFormer and RMSNorm [79, 80] also leverages this setting (up to a global constant). Our analysis can also be extended to standard LayerNorm [81], which also subtracts the mean of  $\pmb{x}$ , while [80] shows that mean subtraction may not affect the empirical results much. LLaMA [82] also uses RMSNorm. Empirically  $\hat{\pmb{u}}_T$  (or  $W_V\hat{\pmb{u}}_T$ ) is normalized (instead of  $X^\top \pmb{b}_T$ ) and here we use an approximation to facilitate analysis: when the token embedding  $\{\pmb{u}_m\}$  are approximately orthogonal to each other, then  $\|U^\top \pmb{x}\|_2 \approx \| \pmb{x}\|_2$  and thus  $\tilde{\pmb{u}}_T \approx \mathrm{LN}(\hat{\pmb{u}}_T)$ .
Objective. We maximize the likelihood of predicted  $(T + 1)$ -th token using cross entropy loss:
$$
\max  _ {W _ {K}, W _ {Q}, W _ {V}, U} J := \mathbb {E} _ {\mathcal {D}} \left[ \boldsymbol {u} _ {x _ {T + 1}} ^ {\top} W _ {V} \tilde {\boldsymbol {u}} _ {T} - \log \sum_ {l} \exp \left(\boldsymbol {u} _ {l} ^ {\top} W _ {V} \tilde {\boldsymbol {u}} _ {T}\right) \right] \tag {2}
$$
For simplicity, we consider single-head attention setting, and multiple-head attention can be regarded as single-head setting with simultaneous different initializations (see Sec. 4). We call  $x_{T} = m$  as the query token of the sequence, and  $x_{T + 1} = n$  as the next token to be predicted. Other tokens  $x_{t}$  ( $1 \leq t \leq T - 1$ ) that are encoded in  $X$  are called contextual tokens. Both the contextual and query tokens can take values from 1 to  $M$  (i.e.,  $m \in [M]$ ) and next token takes the value from 1 to  $K$  (i.e.,  $n \in [K]$ ) where  $K \leq M$ . Fig. 1(a) shows the overall setting. For an overview of the notation used in the paper, please check Tbl. 1 in the Appendix.
3.1 Reparameterization
Instead of studying the dynamics with respect to the parameters of token embedding  $U$ , key, value and query projection matrices  $W_{K}$ ,  $W_{Q}$  and  $W_{V}$ , we study the dynamics of two pairwise token relation matrices  $Y \coloneqq U W_{V}^{\top} U^{\top} \in \mathbb{R}^{M \times M}$  and  $Z \coloneqq U W_{Q} W_{K}^{\top} U^{\top} / \sqrt{d} \in \mathbb{R}^{M \times M}$ . Intuitively, entries of  $Y$  and  $Z$  store the "logits" of pairs of tokens. We regard the empirical parameterization using  $U$ ,  $W_{K}$ ,  $W_{Q}$  and  $W_{V}$  as a specific way of parametrization of  $Y$  and  $Z$ , in order to reduce the number of parameters to be estimated. Previous work also leverage similar parameterization for self-attention layers [27, 46].
For real-world applications, the number of tokens  $M$  can be huge (e.g., the vocabulary size  $M = 50,272$  in OPT-175B [83] and  $M = 32,000$  in LLaMA [82]) and directly optimizing  $Y$  and  $Z$  would be prohibitive. However, as we will show in this work, from the theoretical perspective, treating  $Y$  and  $Z$  as independent variables has some unique advantages and leads to useful insights.
Figure 1: Overall of our setting. (a) A sequence with contextual tokens  $\{x_{1},\ldots ,x_{T - 1}\}$  and query token  $x_{T}$  is fed into 1-layer transformer (self-attention, normalization and decoding) to predict the next token  $x_{T + 1}$ . (b) The definition of sequence classes (Sec. 3.2). A sequence class specifies the conditional probability  $\mathbb{P}(l|m,n)$  of the contextual tokens, given the query token  $x_{T} = m$  and the next token  $x_{T + 1} = n$ . For simplicity, we consider the case that the query token is determined by the next token:  $x_{T} = \psi (x_{T + 1})$  (and thus  $\mathbb{P}(l|m,n) = \mathbb{P}(l|n)$ ) while the same query token  $m$  may correspond to multiple next tokens (i.e.,  $\psi^{-1}(m)$  is not unique). We study two kinds of tokens: common tokens (CT) with  $\mathbb{P}(l|n) > 0$  for multiple sequence class  $n$ , and distinct tokens (DT) with  $\mathbb{P}(l|n) > 0$  for a single sequence class  $n$  only.
Lemma 1 (Dynamics of 1-layer Transformer). The gradient dynamics of Eqn. 2 with batchsize 1 is:
$$
\dot {Y} = \eta_ {Y} \mathrm {L N} (X ^ {\top} \boldsymbol {b} _ {T}) (\boldsymbol {x} _ {T + 1} - \boldsymbol {\alpha}) ^ {\top}, \quad \dot {Z} = \eta_ {Z} \boldsymbol {x} _ {T} (\boldsymbol {x} _ {T + 1} - \boldsymbol {\alpha}) ^ {\top} Y ^ {\top} \frac {P _ {X ^ {\top} \boldsymbol {b} _ {T}} ^ {\perp}}{\| X ^ {\top} \boldsymbol {b} _ {T} \| _ {2}} X ^ {\top} \operatorname {d i a g} (\boldsymbol {b} _ {T}) X \tag {3}
$$
Here  $P_{\pmb{v}}^{\perp} \coloneqq I - \pmb{v}\pmb{v}^{\top} / \| \pmb{v}\|_{2}^{2}$  projects a vector into  $\pmb{v}$ 's orthogonal complementary space,  $\eta_{Y}$  and  $\eta_{Z}$  are the learning rates for the decoder layer  $Y$  and self-attention layer  $Z$ ,  $\alpha \coloneqq [\alpha_{1},\dots,\alpha_{M}]^{\top} \in \mathbb{R}^{M}$  and  $\alpha_{m} \coloneqq \exp(Y^{\top}\mathrm{LN}(X^{\top}\pmb{b}_{T})) / \mathbf{1}^{\top}\exp(Y^{\top}\mathrm{LN}(X^{\top}\pmb{b}_{T}))$ .
Please check Appendix C for the proof. We consider  $Y(0) = Z(0) = 0$  as initial condition. This is reasonable since empirically  $Y$  and  $Z$  are initialized by inner product of  $d$ -dimensional vectors whose components are independently drawn by i.i.d Gaussian. This initial condition is also more realistic than [27] that assumes dominant initialization in diagonal elements. Since  $(\pmb{x}_{T+1} - \pmb{\alpha})^{\top} \mathbf{1} = 0$  and  $P_{X^{\top} \pmb{b}_T}^{\perp} X^{\top} \mathrm{diag}(\pmb{b}_T) X \mathbf{1} = 0$ , we have  $\dot{Y} \mathbf{1} = \dot{Z} \mathbf{1} = 0$  and summation of rows of  $Z(t)$  and  $Y(t)$  remains zero. Since  $\pmb{x}_T$  is a one-hot column vector, the update of  $Z = [z_1, z_2, \dots, z_M]^{\top}$  is done per row:
$$
\dot {\boldsymbol {z}} _ {m} = \eta_ {Z} X ^ {\top} [ i ] \operatorname {d i a g} \left(\boldsymbol {b} _ {T} [ i ]\right) X [ i ] \frac {P _ {X ^ {\top} [ i ] \boldsymbol {b} _ {T} [ i ]} ^ {\perp}}{\| X ^ {\top} [ i ] \boldsymbol {b} _ {T} [ i ] \| _ {2}} Y \left(\boldsymbol {x} _ {T + 1} [ i ] - \boldsymbol {\alpha} [ i ]\right) \tag {4}
$$
where  $m = x_{T}[i]$  is the query token for sample  $i$ ,  $\mathbf{z}_m$  is the  $m$ -th row of  $Z$  and  $\dot{\mathbf{z}}_{m^{\prime}} = 0$  for row  $m^{\prime} \neq m = x_{T}[i]$ . Note that if  $x_{T}[i] = m$ , then  $b_{T}[i]$  is a function of  $\mathbf{z}_m$  only (but not a function of any other  $\mathbf{z}_{m^{\prime}}$ ). Here we explicitly write down the current sample index  $i$ , since batchsize is 1.
3.2 Data Generation
Next we specify a data generation model (Fig. 1(b)), named sequence class, for our analysis.
Sequence Class. We regard the input data as a mixture of multiple sequence classes. Each sequence class is characterized by a triple  $s_{m,n} \coloneqq (\mathbb{P}(l|m,n), m, n)$ . To generate a sequence instance from the class, we first set  $x_{T} = m$  and  $x_{T+1} = n$ , and then generate the contextual tokens with conditional probability  $\mathbb{P}(l|m,n)$ . Let  $\operatorname{supp}(m,n)$  be the subset of token  $l$  with  $\mathbb{P}(l|m,n) > 0$ .
In this work, we consider the case that given a next token  $x_{T + 1} = n$ , the corresponding sequence always ends with a specific query token  $x_{T} = m \eqqcolon \psi(n)$ . This means that we could index sequence class with next token  $x_{T + 1} = n$  alone:  $s_n \coloneqq (\mathbb{P}(l|\psi(n), n), \psi(n), n)$ ,  $\mathbb{P}(l|m, n) = \mathbb{P}(l|n)$  and  $\operatorname{supp}(n) \coloneqq \operatorname{supp}(\psi(n), n)$ .
Note that  $|\psi^{-1}(m)| = 1$  means that the occurrence of token  $m$  alone decides next token  $n$  to be predicted, regardless of other tokens in the sequence, which is a trivial case. When  $|\psi^{-1}(m)| \geq 2$ , the same query token  $m$ , combined with other token  $l$  in the sequence with non-zero probability  $\mathbb{P}(l|m,n) > 0$ , determine the next token.
Overlapping sequence class. Two sequence classes  $s_n$  and  $s_{n'}$  overlap if  $\operatorname{supp}(n) \cap \operatorname{supp}(n') \neq \emptyset$ .
(a) Initialization
(b) Common Token Suppression
Figure 2: Overview of the training dynamics of self-attention map. Here  $\tilde{c}_{l|m,n} \coloneqq \mathbb{P}(l|m,n) \exp(z_{ml})$  is the un-normalized attention score (Eqn. 5). (a) Initialization stage.  $z_{ml}(0) = 0$  and  $\tilde{c}_{l|m,n} = \mathbb{P}(l|m,n)$ . Distinct tokens (Sec. 3.2) shown in blue, common tokens in yellow. (b) Common tokens (CT) are suppressed ( $\tilde{z}_{ml} < 0$ , Theorem 2). (c) Winners-take-all stage. Distinct tokens (DT) with large initial value  $\tilde{c}_{l|m,n}(0)$  start to dominate the attention map (Sec. 5, Theorem 3). (d) Once passing the phase transition, i.e.,  $t \geq t_0 = O(K \ln M / \eta_Y)$ , attention appears (almost) frozen (Sec. 6) and token composition is fixed in the self-attention layer.
(c) Winners-take-all
(d) Attention frozen
(Global) distinct and common tokens. Let  $\Omega(l) \coloneqq \{n : \mathbb{P}(l|n) > 0\}$  be the subset of next tokens that co-occur with contextual token  $l$ . We now can identify two kinds of tokens: the distinct token  $l$  which has  $|\Omega(l)| = 1$  and the common token  $l$  with  $|\Omega(l)| > 1$ . Intuitively, this means that there exists one common token  $l$  so that both  $\mathbb{P}(l|n)$  and  $\mathbb{P}(l|n')$  are strictly positive, e.g., common words like 'the', 'this', 'which' that appear in many sequence classes. In Sec. 5, we will see how these two type of contextual tokens behave very differently when self-attention layer is involved in the training: distinct tokens tend to be paid attention while common tokens tend to be ignored.
3.3 Assumptions
To make our analysis easier, we make the following assumptions:
Assumption 1. We consider (a) no positional encoding, (b) The input sequence is long  $(T\to +\infty)$  and (c) The decoder layer learns much faster than the self-attention layer (i.e.,  $\eta_Y\gg \eta_Z$ ).
Assumption 1(a) suggests that the model is (almost) permutation-invariant. Given the next token to predict  $x_{T + 1} = n$  and the query token  $x_{T} = m$  acted as query, the remaining tokens in the sequence may shuffle. Assumption 1(b) indicates that the frequency of a token  $l$  in the sequence approaches its conditional probability  $\mathbb{P}(l|m,n)\coloneqq \mathbb{P}(l|x_T = m,x_{T + 1} = n)$ .
Note that the assumptions are comparable with or even weaker than previous works, e.g., [27] analyzes positional attention with symmetric initialization, without considering input data and [28] models the data distribution as discriminative/non-discriminative patterns, similar to our common/distinct tokens. Empirically, NoPE [84] shows that decoder-only Transformer models without positional encoding still works decently, justifying that Assumption 1(a) is reasonable.
Given the event  $\{x_{T} = m, x_{T + 1} = n\}$ , suppose for token  $l$ , the conditional probability that it appears in the sequence is  $\mathbb{P}(l|m,n)$ . Then for very long sequence  $T \to +\infty$ , in expectation the number of token  $l$  appears in a sequence of length  $T$  approaches  $T\mathbb{P}(l|m,n)$ . Therefore the per-token self-attention weight  $c_{l|m,n}$  is computed as:
$$
c _ {l \mid m, n} := \frac {T \mathbb {P} (l \mid m , n) \exp \left(z _ {m l}\right)}{\sum_ {l ^ {\prime}} T \mathbb {P} \left(l ^ {\prime} \mid m , n\right) \exp \left(z _ {m l ^ {\prime}}\right)} = \frac {\mathbb {P} (l \mid m , n) \exp \left(z _ {m l}\right)}{\sum_ {l ^ {\prime}} \mathbb {P} \left(l ^ {\prime} \mid m , n\right) \exp \left(z _ {m l ^ {\prime}}\right)} =: \frac {\tilde {c} _ {l \mid m , n}}{\sum_ {l ^ {\prime}} \tilde {c} _ {l ^ {\prime} \mid m , n}} \tag {5}
$$
Here  $z_{ml}$  is  $\mathbf{z}_m$ 's  $l$ -th entry and  $\tilde{c}_{l|m,n} \coloneqq \mathbb{P}(l|m,n)\exp (z_{ml})$  is un-normalized attention score.
Lemma 2. Given the event  $\{x_{T} = m, x_{T + 1} = n\}$ , when  $T \to +\infty$ , we have
$$
X ^ {\top} \boldsymbol {b} _ {T} \rightarrow \boldsymbol {c} _ {m, n}, \quad X ^ {\top} \operatorname {d i a g} (\boldsymbol {b} _ {T}) X \rightarrow \operatorname {d i a g} (\boldsymbol {c} _ {m, n}) \tag {6}
$$
where  $\mathbf{c}_{m,n} = [c_{1|m,n}, c_{2|m,n}, \ldots, c_{M|m,n}]^{\top} \in \mathbb{R}^{M}$ . Note that  $\mathbf{c}_{m,n}^{\top} \mathbf{1} = 1$ .
By the data generation process (Sec. 3.2), given the next token  $x_{T + 1} = n$ , the query token  $x_{T} = m$  is uniquely determined. In the following, we just use  $c_{n}$  to represent  $c_{m,n}$  (and similar for  $\tilde{c}_n$ ).
4 Dynamics of  $Y$
We first study the dynamics of  $Y$ . From Assumption 1(c),  $Y$  learns much faster and we can treat the lower layer output (i.e.,  $X^{\top}b_{T}$ ) as constant. From Lemma 2, when the sequence is long, we know
given the next token  $x_{T + 1} = n$ ,  $X^{\top}\pmb{b}_{T}$  becomes fixed. Therefore, the dynamics of  $Y$  becomes:
$$
\dot {Y} = \eta_ {Y} \boldsymbol {f} _ {n} \left(\boldsymbol {e} _ {n} - \boldsymbol {\alpha} _ {n}\right) ^ {\top}, \quad \boldsymbol {\alpha} _ {n} = \frac {\exp \left(Y ^ {\top} \boldsymbol {f} _ {n}\right)}{\mathbf {1} ^ {\top} \exp \left(Y ^ {\top} \boldsymbol {f} _ {n}\right)} \tag {7}
$$