Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers
Authors: Siyu Chen
Heejune Sheen
Tianhao Wang
Zhuoran Yang
Venue: NeurIPS 2024 poster (2024)
Keywords: transformer
in context learning
Markov chain
n-gram model
feature learning
training dynamics
induction head
Primary area: optimization
TL;DR: We theoretically prove that a two-layer transformer model learns to perform "induction head" after training on n-gram Markovian data

Abstract: In-context learning (ICL) is a cornerstone of large language model (LLM) functionality, yet its theoretical foundations remain elusive due to the complexity of transformer architectures. In particular, most existing work only theoretically explains how the attention mechanism facilitates ICL under certain data models. It remains unclear how the other building blocks of the transformer contribute to ICL. To address this question, we study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data, where each token in the Markov chain statistically depends on the previous n tokens. 
We analyze a sophisticated transformer model featuring relative positional embedding, multi-head softmax attention,  and a feed-forward layer with normalization. 
We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model that performs a generalized version of the "induction head" mechanism with a learned feature, resulting from the congruous contribution of all the building blocks. 
Specifically, the first attention layer acts as a copier, copying past tokens within a given window to each position, and the feed-forward network with normalization acts as a selector that generates a feature vector by only looking at informationally relevant parents from the window. 
Finally, the second attention layer is a classifier that
compares these features with the feature at the output position, and uses the resulting similarity scores to generate the desired output. Our theory is further validated by simulation experiments.

Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers
Siyu Chen"
Heejune Sheen"
Department of Statistics and Data Science, Yale University
Department Statistics and Data Science,
Yale University www.eimwee..eeeeeeee...mmm edu
siyu. www.. www.scccaeccccccccmmmmm
Tianhao Wang
Zhuoran Yang
Toyota Technological Institute at Chicago tianhao. wang@ttic.
Department of Statistics and Data Science, Yale University
zhuoran zhuoran. wwn.gyagggyaeeeee
Abstract
In-context learning (ICL) is cornerstone of large language model (LLM) fune- tionality, yet its theoretical foundations remain elusive due to the complexity transformer architectures. In particular, most existing work only theoretically ex- plains how the attention mechanism facilitates ICL under certain data models. It remains unclear how the other building blocks of the transformer contribute to ICL. To address this question, study how a two-attention-layer transformer is trained to perform ICL on и-gram Markov chain data, where each token in the Markov chain statistically depends on the previous tokens. We analyze sophisticated transformer model featuring relative positional embedding, multi-head softmax attention, and a feed-forward layer with normalization. We prove that the gradient flow with respect cross-entropy ICL loss converges to a limiting model that performs a generalized version of the "induction head" mechanism with a learned feature, resulting from the congruous contribution of all the building blocks. In the limiting model, the first attention layer acts as a copier, copying past tokens within given window to each position, and the ferd-forward network with normalization acts as a selector that generates feature vector by only looking at informationally relevant parents from the window. Finally, the second attention layer is a classifier that compares these features with the feature the output position, and uses the resulting similarity scores to generate the desired output. Our theory is further validated by simulation experiments.
Introduction
In-context learning (ICL) (Brown et al., 2020) has emerged crucial aspect of large language model (LLMM) (Radford et al., 2019; Brown al., 2020; Achiam et al., 2023; Anthropic, 2023; Team et al., 2023) functionality, enabling pre-trained LLMS solve user-specified tasks during inference without updating model parameters. In ICL, pre-trained LLM, typically a transformer, receives prompts containing a few demonstration examples sampled from task-specific distribution and produces the desired output for that task. This capability is noteworthy because the tasks addressed during the ICL might not be part of the original training data set. The success ICL requires the LLM to perform certain learning processes during inference.
"qqua contribution
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

Although many previous works aim to demystify ICL from either empirical or theoretical perspec- tives, the theoretical foundations of ICL remain elusive. This is primarily due to the complexity of transformer architectures, which integrate token and position embeddings, multiple layers of multi- head softmax attention, layer normalization, and feedforward neural networks. When it comes to understanding how the ICL ability emerges in transformers after training, existing works often focus on simplified models, such as linear attention mechanisms or single-layer transformers (Von Oswald et al., 2023), and ICL tasks are typically confined to linear regression (Akyürek et al., 2023). This leaves gap in understanding how full-fledged transformer architectures facilitate ICL of more complex tasks, especially when latent causal structures exist among the tokens in sequence.
In this paper, our aim is to narrow this gap by studying how two-attention-layer transformer is trained to perform ICL of an n-gram Markov chain model, where each token in the Markov chain statistically depends on the tokens before it, known as the parent set. Specifically, we consider a transformer model with relative positional embedding (RPE) (He et al., 2020), multi-head softmax attention, and a feed-forward network (FFN) layer with normalization. We employ such transformer model to predict the (L + 1)-th token of an n-gram Markov chain, with the first L tokens given as the prompt, where L + 1 is the sequence length. Here the L-token sequence is sampled from a random Markov chain model, where a random transition kernel obeying the n- n-gram Markov property is used to generate sequences. The token sequence is fed into the transformer model, which outputs a probability distribution over the vocabulary set to predict the (L + 1)-th token. To train the transformer model, we sample token sequences from these random Markov chain models and minimize the cross-entropy loss between the predicted token distribution and the true token distribution.
Under this setting, we aim to answer the following three questions: (i) Does the gradient flow with respect to the cross-entrope loss converge during training? (ii) yes, how does the limiting model perform ICL? (iii) How do the building blocks of the transformer model contribute to ICL?
Main Results. We provide an affirmative answer to the Question (i) by proving that the gradient flow converges during training. In particular, we identify three phases of training dynamics: in the first stage, FFN learns the potential parent set; in the second stage, each attention head of the first multi-head softmax attention layer learns to focus on a single parent token selected by FFN; and in the final stage, the parameter of the second attention layer increases, and the transformer approaches the limiting model. Moreover, for Questions (ii) and (iii), we show that the limiting model performs a specialized form of exponential kernel regression, dubbed "generalized induction head", which requires the congruous contribution of all the building blocks. Specifically, the first attention layer acts as copier, copying past tokens within a given window to each position. The FFN layer acts as selector that generates feature vector by only looking at informationally relevant parents from the window according to a modified x" -mmutal information. Finally, the second attention layer an exponential kernel classifier that compares the features at each position with those created for the output position L + 1. and uses the resulting similarity scores to generate the desired output. When specialized to the case where n = 1. the limiting model selects the true parent token and implements the induction head mechanism (Elhage al., 2021). In this case, we recover the theory in Nichani et al. (2024). Our theory is complemented by numerical experiments, which validate the three-phase training dynamics and mechanism of generalized induction head.
To best knowledge, our work is the first provide a comprehensive understanding of how ICL is empowered by a collaboration of different building blocks in a transformer model. In particular, we identify the pivotal roles played by RPE in the copier component, the FFN layer with normalization in the selector component, and attention in the classifier component. We believe our work will shed light on the theoretical understanding of ICL for more complicated tasks.
Related Works Our work adds to the rapidly growing literature on understanding in-context learning by transformers. We defer an in-depth discussion on related works in Appendix $A due to the page limit.
Roadmap. The rest of the paper is organized as follows: We introduce the problem setup of ICL of Markov chains in $2. Then in $3, we present the main theoretical results and related discussions. A proof sketch is provided in SD. Finally, we present corresponding experiment results in $B, and the detailed proofs are deferred to the Appendix.
Notation. We denote by e1, Ea the standard basis vectors in I2d and by the all-one vector in Idd. We denote by (-) the softmax function such that the i-th coordinate of (z)) is 0,(x)
2

---

exp(zi)/Ei= exp(xi) for E R.", By default, the softmax operation will always be applied row- wise. For any integer > O, we denote [n)
.......n... For a vector E RM we denote by the i-th entry of and w-; the (M + i)-th entry of w for positive integer E [M For matrix W. we denote by ji, the entry at the i-th row and j-th column of W. For two vectors u and v, write u/u the vector obtained by taking element-wis division between and u. We denote by a V and a A the maximum and minimum of a and b, respectively. We denote by Tast the sequence {Ia, [T -....-....:---............... J. For class X' we denote by ( the space of probability measures over X. We use the standard big O notation throughout the paper.
2 Problem Setup: In-Context Learning of Markov Chains
In this section, we present the details of the problem setting. particular, we first introduce the statistical problem of ICL of n-gram Markov chains in $2.1 and then lay out the details of the transformer model in $2.2.
2.1 In-Context Learning and n-Gram Markov Chains
We study how autoregressive transformers are trained to perform in-context learning (ICL). A pre- trained transformer can be viewed as conditional distribution ftr(-] prompt) over a finite vocabulary set X, where prompt is a sequence of tokens in X. We consider an in-context unsupervised learning problem where the pre-trained transformer Jug is used to predict the (L + 1)-th token #L+1 with the first tokens being the prompt. Here a fixed number and the joint distribution of the sequence 11:(1+1) is sampled from a random n-gram Markov chain. In other words. with (1:(++1) sampled from some distribution, we evaluate how well fir (1LE)) predicts the distribution of #L+1-
n-Gram Markov Chains. We assume the data comes from a mixture of n-gram Markov chain model, denoted by tuple (X,
pa, P, Po)) where X' the state space and pa = (-71....: -Tn) the parent set with positive integers ri r2 < ... < Tx.
X1-2
x
That is, for each I > Fn. 2y only statistically depends on (Ii-r...... I(-r,), which denoted by X Xpa(!) and referred to as the parent tokens of zt. We let
Figure 1: A two-gram Markov chain with par- ent set pa (-1,-2)
|X' denote the vocabulary size, Moreover, P a probability distribution over the set Markov transition kernels respecting the parent structure specified by pa, and poo is the joint distribution of the first rn tokens Fi.r, Note that the size of the parent n can be smaller than or equal Tn- Thus, the sequence (1:((+1) is generated as follows: (i) sample initial ra tokens (æ), Tr,)) Alo. (ii) sample a random transition kernel π P. where π: X + A and (iii) sample token Ii x-- (xI])att)) for Tn I,...,L. L 1. See Figure I for illustration of the generating model of #1:(L+1)-
Cross-Entropy Loss. When F1:(L++1) is generated, #1.L is fed into the transformer fur to predict F1+1. To assess the performance of ICI., we adopt the population cross-entropy (CE) loss
C(fe)) -EE P (1og(fe((E1+1 T1:2) + +c]]
(2.1)
where > 0 is small constant introduced for numerical stability and in the sequel we will take O(L-1/2). Here, the expectation is taken with respect to the joint distribution of (1:(++1) (including the randomness of ~ P). When setting = 0, we note that minimizing this cross- entropy loss is equivalent to minimizing the KL divergence
EgwP. P. [L(( Xpe(L+1)) || See( [1.)))]-
As remark, we also relax condition in Nichani al. (2024) where the last token øy has to be resampled from uniform distribution. In addition, our analysis can also be extended to sequential CE loss, which corresponds predicting every token in the sequence given the past rather than just the last token 11-1. This is closer to the training paradigm used in practice (Brown et al., 2020). See $C.4 for further discussion the sequential CE loss.

---

2.2 A Two-Layer Transformer Model
We consider a class of (wo-attention-llyer transformer model, denoted by TF(M, H, D), which incorporates Relative Positional Embedding (RPE) (He et al., 2020), Multi-Heate Attention (MHA) network (FFN) with normalization. Here M is integer
(Vaswani et al., 2017), and that specifies the window size of RPE, is the number of heads in the first attention layer, is the vocabulary size, and D is an integer that controls the complexity of FFN. The details of (F(M, H, d, D) are as follows.
Token Embedding, Input and Output. Note that each token takes values in N' with d = - We embed the tokens into one-hot vectors R4, and thus we can identify X as the canonical basis in Fd i..., X' = .e...... c.). ) A transformer model can be viewed as mapping from [[+++)) to (Y)) In particular, given the input sequence F1:L, we denote = (I1, ................................ (.ll,.,,,,.))).. xL) Rxxd, and we append zero vector E IRd to the sequence, and define X (xI...., I1, 0) E (R(++)xdd The transformer takes X as input and outputs a probability distribution over X.
Relative Positional Embedding. In each head of the first attention layer, we adopt RPE to incorporate positional information. Specifically, RPE is parameterized by vector w = (W-M,-- -....... w- w-1)² E R.M. and it assigns a scalar value W Wr(i,)) to a pair of positions (i. satisfying
10
w.M
Wr(1, = w i - j (1....,M), M): Wr(i,5) = if j or j j- > M.
In other words, illustrated in Figure 2, the i-th token only attends to tokens with indices {i {i-1,.. 1,..., M). referred to as the length-M M window of the i-th token, and the trainable vector w determines the value of positional embedding. Here, we use -k to index the last k-th posi- tion.
[b)
Figure 2: Illustration of the relation- ship between RPE vector (h()) and cor- responding matrix W!!h)
p
The First Attention Layer. The input sequence is processed by the first attention layer with H parallel heads. In all heads, we discard the token information and only use RPE to compute the attention score. Specifically, each attention head h maps X into sequence in Rd with length L + 1, (h) E (L++)idd For any E [I. +
denoted by (hh) (u) 241)
1}. oz
(h) is computed via
L
(h) exp(WV (1,5)) xj JJ (W (M) (1,)) 2j Et exp (W/ p (1,
L
(b)
Σ
(2.2)
That is, we use the RPE parameter w to construct a weighted sum over the input sequence at each position 1 E [T++1] 1J. Here w((h) is the RPE matrix of the h-th head.
Feed-Forward Network with Normalization. Following the first attention layer, we concatenate the outputs of the H attention heads and define V (V(2)... ,..,, V(H) IR (B++1)xxHd Here we abuse the notation and write V = (e)..... UL+1) T, i.e., each t is the 1-th row of For any vector Hd, can split it into (o()) .....
(H)) (())T) where each block e(h)) Rd. For embedding dimension de, each vector of V is passed through an FFN o(-) Rdd - Rd,, which specifies a polynomial kernel such that for any v. v E Hdd we have
(o(o), ((())
Σ
II (w()) (r(a)) RES
(2.3)
SE(M SD
Here, the low-degree parent set [H]DD == ( [H] [S] S D) contains all subsets of [H] with cardinality at most D, and (cs S [M] are the corresponding trainable parameters of Φ(-). Therefore, the FFN <b-) specifies kernel on the output of the multihead attention which induces a special inner product structure. While (2.3) characterizes o(-) implicitly, we provide an explicit construction of 6(-) in Lemma C. vector valued mapping whose entries are monomials of the

---

input's entries. Moreover, the complexity of o(-) is controlled by the maximum degree D, which also influences the embedding dimension de as we show in the construction.
Furthermore, to control the magnitude of the FFN outputs, we normalize Φ(+) by letting oi ((n)/vCD for all + I]. where define CD EsejHig S Such normalization scheme is motivated by the standard layer normalization (Ba et al., 2016) in transformer architectures. To motivate the use of vCD as the normalization, consider special case where the positional embeddings, after the softmax function, produce attention weights that are close to one-hot for each
(h) in (2.2) is equal some token in T1:L. As result, each consists of H tokens and head. Then U/
[(0())))22
Σ ds SE(H|ED
1o (1) 1
(1)) VCD.
RES
Thus, u, is roughly equivalent to the output of the layer normalization of ((ux)//l((zz))l) (without trainable parameters). Although our theoretical analysis and simulations focus on this simplified version of layer normalization, our additional experiments in $B.2 demonstrate that it aligns well with the performance of the actual layer normalization.
The Second Attention Layer. The normalized vector sequence U (u,,..., U1+1) T and the original sequence X are then fed into the second attention layer to generate the final output. In particular, U2+1 is used as the query to compare with the keys {uM+1,... , WL J. and the resulting attention scores are used aggregate
values T(M+1):L- This attention layer has a single head and a scalar trainable parameter a. We let UL = (u1...., "z) E Lxxds and denote by Mask(-) .) the mask that sets every entry of the first M rows of matrix to be -oo. The final output is given by
exp(a u]+ 2z
V Σ
a,(a.un u Mask(U Ij
Σ
(2.4)
Et exp(a u "oo.o[[[1"" (())
M=M+1
Note that the softmax function in (2.4) yields a probability distribution over [L] and that Ti.r is sequence of one-hot vectors. Thus y in (2.4) is probability distribution over X. The mask operator included here just to simplify our analysis while the experiments we are not using the mask.
In summary, given the input X [(++1)xd the matrix form, our transformer model TF,M, H, D) consecutively applies the following operations:
First Attention:
(h) = (
E R((-1)xd Vh E [H];
Concatenate:
V = [V()) (y()....,((I)) v(H)
E RR(-I)xHd,
(2.5)
FFN & Normalize:
= VVV)/VCD
E ((4-1)xde
Second Attention:
y = oaaa (a uT (.s
E Rlxd
The trainable parameters the above transformer model are denoted by
θ {a, (w 1 ..... w (h)
heem, {cs S [[]g
We remark that the transformer model in (2.5) is known as a disentangled transformer (Friedman et al., 2024), which is version of the transformer model that is One thing be noted is that there a residual connection that directly copies X the output of the FFN & Normalize block, which gives us [U, X). and the second attention layer will treat the copied X as the value in the attention mechanism. We omit the residual connection in the above paradigm for notation simplicity. As shown in Nichani et al. (2024), any standard transformer model can be expressed as a disentangled transformer by specializing the attention weights to allow feature concatenation.
amenable for theoretical analysis.
Our goal is investigate whether the transformer model TF(M, H, d, D) perform ICL n-gram Markov chains and further, whether such capability can be learned from data with common training algorithms like gradient descent.
5

---

3 Theoretical Results
In this section, present the theoretical results. We show in $3.1 and $C.1 that there exists transformer in TF(M, H, d, D) that implements a generalized "induction head" mechanism (Oisson et al., 2022) with learned feature, which serves as natural algorithm for learning n-gram Markov chains. Then $3.2 we prove that the gradient flow in (3.4) finds such a desired model asymptotically.
3.1 Generalized Induction Head Mechanism for Learning n-Gram Markov Chains
Recall that we define the mixture of n-gram Markov chain model (X, (X,pa, Mo) in $2.1, where P is a distribution over the Markov transition kernels. For regularity, we assume existence of a unique stationary distribution for any π E supp(P), where a rigorous statement deferred to Assumption 3.5. We also assume the window size M > Tn- For any n--ram Markov chain with transition kernel P. we let μº E ((x*+11 denote the stationary distribution of the Markov chain over a window of size M + 1. Here we use {=/)>= to denote a random sequence of tokens generated by the Markov chain. Then jt" denotes the joint distribution of a block of M + 1 tokens (#-M, 21-1; 21) under the stationary distribution of π, where > M is an integer.
In the following, we introduce a generalized induction head (GIH) estimator for the task of predicting FL+1 given T1:L, which is based the following simple idea: FL+11 should be similar to a previous token zι if their parents are similar. As the parent set pa is unknown, GIH adopts an information- theoretic criterion to select a subset of previous tokens as proxy of the parents. Specifically, GIH uses modified version of x2 2 information, which is defined as follows.
Definition 3.1 (Modified X2-Mutual Information). We take length- (M + 1) windows (2-M,...- 11-1; 21) for some I. > M and suppose the sequence is sampled from stationary dis- tribution n² with ~ ~P. Let Z = (21-M...... ...1). For any subset S _ [M), we use Z to denote the subvector of % containing entries of the form zi- Vs -. For instance, suppose S = (2, then Z-s = (21-5; A1-2) The modified x" -²-mutual information for S is defined as
{Σ EEX
1 (====e) ---2--)) [=(3)==12-5) -
Tx-(S)= E Bx,D,
11" (2-s)
(3.1)
where p= = ()=s) is the conditional distribution of induced by nº' given the partial history Z-s, and μº (Z-s),, 166 (z) are the marginal distributions of Z-s and z under (z, Z) ~ ye.
Intuitively, (S) is modified from the vanilla 2mmutual information Mx2 between two random variables (Polyanskiy and Wu, 2024) and quantifies how much information the partial history Z_s contains about In particular, we incorporate an additional μ (Z-s) term that decreases with the growing size of S. To see the rationality, we first introduce GIH estimator based on the modified -2-mutull information.
Definition 3.2 (Generalized Induction Head), GIH estimator with window size M E N, feature size D E N is denoted by GIH(-; D), which maps I1LL to distribution over X. We let S be the information optimal subset (referred to as the "information set" in the sequel") of [M] with size no more than D that maximizes the modified x-mutual information Tx)(.) defined in (3.1). That is, we define the information set S as
S argmaasseMM (S)
(3.2)
Then (1111 M, D) outputs
N Et-M+1 21 1(Xi-s = = fN 2N1, 1,
5" = (L- M)-1 Et=m+1 Nt: otherwise.
(3.3)
Here, we define Xi-s- as the set (2i : 8 S*) and = -EL 1(Xi-s- = XL+1-S.).
Note that S* defined in (3.2) depends on the choices of M and D and serves as a proxy of the unknown parent set pa based on (-) defined in (3.1). In a nutshell, the GIH estimator checks
With slight abuse of notation, we also call Xi-s- = (x1- 8 E S*) the information set of the I-th token 21-

---

whether the partial histories of Xi-s- and X2+1-S- match and aggregate all the tokens 2 that have a matching partial history to predict L As a remark, using the modified x 2MI as the information criterion rules out redundancy in the information set S in the following sense:
S cannot be a superset of the true parents. Note that if S is a superset of the true parent set, by the Markov property, and Z-s are conditionally independent given the true parents Zpa. Thus, maximizing the vanilla x x -mutual information yields multiple maximizers, i.e., all the supersets of the true parent set. However, with the modification in (3.1), any superset yields strictly smaller Ix² compared to the exact parent set, making them suboptimal.
The modified -MI selects informative partial history. Even true parent may bear relatively little information about the target compared to other parents sometimes. Meanwhile, exact match of a larger set of partial history becomes much harder as it tends appear less frequently in the context sequence, leading to poor estimation accuracy for the estimator in (3.3). The modified x2-MI reaches a balance by selecting the informative partial history while penalizing the size of the information set.
The term involving 1= = - can be viewed as the signal part which helps us to find an informative subset S. The term (Z-s) can be viewed as penalty on the model complexity which favors smaller subsets. Thus, the modified x²-MI strikes a balance between these two objectives and enables us find a good proxy S* of pa when L finite. Moreover, when L is sufficiently large, we identify two scenarios in which maximizing (+) yields the true parent set (see $C.7 for details). Moreover, the GIH estimator is generalization of the induction head mechanism (Elhage et al., 2021) to the stochastic setting with multiple parents, where we give the model more flexibility to learn based on a partial history that does not necessarily correspond to the true parent set. As we will show in 1, the GIH mechanism can be implemented by the transformer model.
3.2 Convergence Guarantee of Gradient Flow
In the following, present the convergence guarantee for gradient flow. To simplify the discussion, we consider the case where = M, meaning there are enough heads to implement the GIH mechanism by having each head copy unique parent token from window of size M. Let us first introduce the paradigm of training by gradient flow.
Training Paradigm. Consider training transformer TF H. d. D) in (2.5) with M H to perform ICL on the n-gram Markov chain model introduced in $2.1. Specifically, we define (()) as the population cross-entropy loss in (2.1), where the transformer model für is given by (2.5) with parameter O Ideally, when training the parameter with gradient flow, the dynamics with respect to the loss C(O) is given by:
a,0(t) -VC()(())
(3.4)
We consider three-stage training paradigm where, in each stage. only specific subset of the weights trained by gradient flow. The three stages are outlined in Table I. Specifically, in the first stage, we only train the FFN layer via gradient flow while keeping other weights fixed. We then only train the RPE weights in the first attention layer in the second stage. Finally, we only train the weight in the second attention layer in the last stage, while fixing the rest of the parameters. This training approach is primarily used for analytical convenience; in practice, the entire model can be trained simultaneously. and similar convergence results are reported in $B.2. From theoretical standpoint, we will also justify the three-stage paradigm in the discussion following Theorem 3.6.
Initialization Conditions. Before presenting our main results about how training by gradient flow induces the GIH structure, let us introduce the following assumption on the initialization of the weights. We define the information gap within the D-degree parent set [ as
Aī. Ix²
max SE(H|gn (s).
(3.5)
where we recall that S* defined in (3.2) maximizes the modified x2 mutual information.
Assumption 3.3 (Initialization). We assume that the following holds at initialization:
7

---

Stage
Weights to Train
Description
(cstsE(M)ss in the FFN layer
Ratio cs- cs (t) grows exponentially, learning the low low-degree features with S*,
II
Tul"))eein in the RPE of
1 1- II (o (t))² decays polynomi- ally
the first attention layer,
training each head in S* to be copier, a(t) experiences a two-stage growth, learning the softmax aggregator for GIH,
III a in the weight of the second attention layer
Table 1: Three-stage training paradigm for gradient flow. Here, the "Weights to Train" column indicates the weights updated in each stage, and the "Deseription" column summarizes the corresponding results from Theorem 3.6.
1. For the first attention layer's 's RPE weights, w"h 2 wh + Awfor all h,j E [H] with j & h, where Aw > O is a positive scalar satisfying
Aw 2 log(M - -1) - log
( (1+
-1). -1)
(3.6)
2. The scalar parameter in the second attention layer satisfies a S O(L-0(L)))
The first assumption on the RPE is used to induce the correspondence between parents and heads during the training by slightly breaking the symmetry between different attention heads. The second assumption on the scale of a ensures that the attention probability given by the second attention layer is close to the uniform distribution over [L]. These initialization conditions enable us to derive clean descriptions for the dynamics of the first attention layer and the FFN, shedding light on their respective roles in executing ICI.
We now outline our assumptions on the Markov chain used in the data generation process. Recall that r, is the largest absolute integer in the parent set pa. For any position I, we define the history Z = (2) E........... 1) as the last state and Z' (2) (2)-r++1,.... 21) as the current state. Since the parent of the new token 2 already included Z, Z' is independent of all prior history given Z, forming a Markov chain.
We define Pa as the d' x d transition matrix for this Markov chain, where states are successive Tn tokens. Each row of Pa is indexed by Z' and each column by Z. The matrix element (Z',', Z)
is thus given by
( =) = ( 7pa(1)) 1(2) = Z Z--roti--1..
This means that to transition from Z to Z', all elements of Z' except for z'_1 must match the last Yn tokens of Z. The token z/ is then sampled according to the transition kernel æ and depends only the parent 7 The above definition is in fact independent of the position / as the transition kernel π is the same across all positions. Note that P. is also a stochastic matrix but with zero entries due to the indicator. To proceed, we need the following notion of primitive matrix to state our assumption on Px.
Definition 3.4 (Primitive Matrix). A nonnegative and irreducible square matrix P is called primitive if there exists a positive integer k such that all entries of pk are positive.
We defer more details about the above definition to SC.3. By the celebrated Perron-Frobenius theorem, a stochastic matrix P, is also primitive, then (i) there exists a unique stationary distribution for the Markov chain; (ii) P, has a unique leading eigenvalue equal to 1, and the corresponding eigenvector the stationary distribution. Next, we state the assumptions on the mixture of Markov chains for data generation.
Assumption 3.5 (Markov Chain). For any π E supp(P), we assume that:
1. The transition matrix P= is primitive. In particular, we assume that there exists X < 1 such that the eigenvalue of P, with the second largest magnitude satisfies ()2(Px) < X. Note that (2()x) can be complex-valued.
8

---

2. There exists y > 0 such that the transition kernel satisfies π( Xpa) 2 for any (z, Xpa).
In fact, the second condition <-- |Xpa) > y already ensures that P7 must be primitive, as is required by the first condition. See Corollary F.14 for details. On the high level, the first assumption guarantees unique stationary distribution as well as a fast mixing rate of the Markov chain by ensuring spectral gap for Px. The second assumption implies a lower bound on the probability for any set S [M] under the stationary distribution, i.e., 1r8 (X1-s) 2 7/8 for any I > M. See Corollary F115 for details.
Now are ready to present our main theoretical result training transformers by gradient flow.
Theorem 3.6 (Convergence of Gradient Flow). Suppose Assumption 3.3 and Assumption 3.5 hold. Consider H 2 M. We set & = L--1/2 for the cross-entropy loss and assume is sufficiently large. Then the following holds for the three-stage training of gradient flow:
Stage I: Parent Selection by FEN. Let Cp(t) = ESEEMin cs(t)² and ps- cz.
Then in the first stage with duration t x (D) log L/(a(0)AI,z) ), the ratio cs /cs grows exponentially fast for any S + S*, and S* dominates exponentially fast in the sense that,
1 1 1-ps-(t) S (1 -
(1-ps- (O) exp(((200)) a(0) . pss (0) -Ixa- t),
Vt (0,t).
Stage II: Concentration of The First Attention. Define (th) = o( (()) E RM, and let
mial(t) == minkes. a - (t).
(t). Then in the second stage with duration t2 t = (xa the first layer's attention heads have attention probabilities concentrated on the optimal infor- mation set S* in the sense that for any E [11, t1 + t2),
1- II hEST
(a (10
(0)
S
2 (M - 1)
A
a())
AIxa min (0) (1 (1- 11) exp(Aw) +
Stage III: Growth of The Second Attention. For some constants C1, C2 depending on (P.S*) with 0 < or C2, there exists small constant ö > 0 such that the growth of a(t) exhibits the following two sub-stages: (i) When a(t) < log(cr/8), (8). it holds that da(t) x exa(t); ; (ii) After a(t) has grown such that a(t) 2 log(c2/ /8), then 0,a(t) = 1/a(t) until reaches the value log T./8.
See $D for a proof sketch and SE for the detailed proof. We require that T. is sufficiently large, and the specific conditions for L are deferred to SE.1.
Interpretation of Training Dynamies. We empirically verify Theorem 3.6 by conducting sim-
ulation experiment. In particular, we train a transformer with = M = 3 and D = 2 based on Markov chain data with d = 2, L 100 and pa = {-1, -2). We sample the transition kernel from Dirichlet prior such that S* {1, 2} also matches the parent set. For more details on this simulation, see $B. The results are shown in Figure 3 and align perfectly with Theorem 3.6. From Theorem 3.6, we can interpret the three stages of training dynamics as follows.
In the first stage, the training of FFN parameters learns a selector that selects an informative set S* by realizing the corresponding feature embedding through the polynomial kernel. That is, when t is sufficiently large, we have ps- (t) R 1 and (t()) se for all S a S*. In this case, for any input vectors v, E Hdd the inner product (2.3) reduces to
(o(v),<(v') = cs.
II (o(1) ((1))) BES
That is, FFN only selects the blocks in St as the feature. We observe this phenomenon in the experiment, where we set S* = [1,2]. As shown in Figure 3-(a), it is clear that cs- immediately dominates the rest of cs' within only a few gradient epochs.
In the second stage, we update the parameters of the RPE. This stage turns the first attention layer into copier by establishing the correspondence between the attention heads and the parents in the selected S* That is, each attention head copies a particular parent in
S*. Specifically, when t is sufficiently large, for any h E St, ((h)) (t) = (w((h)) (t)) Rs 1.
(h) Recalling the construction of RPE, this implies that t7 in (2.2) becomes zı-h for all

---

pairs of the form (xt)) to predict LL+1. In particular, when a(t) goes to infinity, the transformer exactly becomes the GIH Stage E: Train
Siage Train wP)
Stage III: Train
1.0
0.8
0.6
0.4
0.0
3000
4000
Epoch
Beginning
Ending
Epoch
(c)
Figure 3: An illustration of the transformer parameters during the three-stage training. We train transformer TF(M 3, H = 3, d 3. D 2) with T. 100, pa {-1, 121.-2..
22. See $B
and Figure 4 for more details of the simulation.
h E S*. As shown in Figure 3-(b), in the experiment, the first two heads initialized towards the first two parents will deterministically copy parents -1 and -2 eventually. The third head stays close to its initial value. This head has a negligible effect on the output because 3 & S* and ps* x 1.
After the first two stages are completed, know that the features constructed approximately satisfy (C.1) up to proportionality factor. Then, in the final training stage, the scalar weight in the second attention layer keeps increasing. Thus, this stage learns exponential kernel classifier as specified in (C.2). When al)) is sufficiently large, the learned transformer is close to classifier that uses covariate-
mechanism given in Definition 3.2. Moreover, we theoretically prove that the increasing trajectory of at)) has two stages, where da(t)/dt is initially large and gradually decays, this is also clearly observed in the experiment. See Figure 3-s(c) for details.
In summary, we theoretically show that the limiting model obtained by three-stage training ap- proximately implements the GIH mechanism. We will prove that the difference between these two estimators is at most (1/8) We defer the formal statement and proof to $E.5. Moreover, as an answer the Question (iii) raised in $1, the different components of the transformer architecture all critical for achieving this: FFN with normalization realizes the selector, the multi-head design of attention supports the copier, and finally, the softmax operation facilitates the exponential kernel classifier. These components work organically as a whole system, yielding the trained transformer's capability of ICL of n-gram Markov chains.
Another takeaway from Theorem 3.6 is strict separation in the growth rate of these three stages. In particular, the convergence rates of the corresponding components of the transformer model in these three stages range from exponentially fast (Stage I), polynomially fast (Stage II), to logarithmically slow (Stage III). With such two exponential separations of convergence rates, we expect that these three stages naturally arise when we simultaneously train the whole model via gradient descent/flow. We empirically verify this argument and the details are deferred $B.2.
In $C.7, we provide more intuitive interpretation of the modified x²-mutu information, which demonstrates a balance of model complexity and information richness.
4 Conclusion and Future Work
In this paper, we have studied the training dynamics of two-attention-laye layer transformer model for learning n-gram Markov chains in an in-context way. Our work opens new directions for developing rigorous understanding of the transformer models, which includes understanding the induction head mechanism with standard FFN layer and investigating the training dynamics beyond a single loop of this induction head mechanism. We defer readers to $C.8 for more discussions.
10

---

Acknowledgement
We acknowledge Shaobo Wang for his help with the experiments. We also thank Jason D. Lee, Alex Damian, and Eshaan Nichani for their helpful discussions. Zhuoran Yang acknowledges the support of NSF under the award DMS-2413243.

---

References
Achiam, J., Adler, S., Agarwal, S., Ahmad, I.,, Akkaya, L., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S. et al. (2023). Gpt-4 technical report. preprint arXivv2303..06 087774.
arXiv
Ahn, K., Cheng, X.. Daneshmand, H. and Sra, S. (2023). Transformers lear implement precondi- tioned gradient descent for in in-oontext learning, arXiv preprint 200.092300.00297.
Ahuja, K., Panwar, M. and Goyal, N. (2023). In-context learning through the bayesian prism. arXiv preprint arXiv:2306.0443
Akyürek, E., Schuurmans, D., Andreas, Ma. T. and Zhou, (2023). What learning algorithm is in-context learning? investigations with linear models, In The Eleventh International Conference Learming Representations.
Alayrac, J.-B., Donahue, J., Luc, P., Miech, A., Barr, Hasson, Y., Lenc, K., Mensch, A., Millican, K., Reynolds, M. et al. (2022). Flamingo: a visual language model for few-shot learning. Advances in neural information processing systems, 35 237 23716-23736.
Anthropic (2023). Model card and evaluations for claude models.
Ba, L., Kiros, J. R. and Hinton, G. E. (2016). Layer normalization. arXiv:1607.064450
arXiv preprint
Bai, Y. Chen, F., Wang, H., Xiong, C. and Mei, S. (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint 2rX6v:2306.04466
Bietti, A., Cabannes, V., Bouchacourt, D.. Jegou, H. and Bottou, L. (2024). Birth of a transformer:
A memory viewpoint. Advances in Neural Information Processing Systems, 36.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A.,
Shyam, P., Sastry, G., Askell, A. et al. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33 1877-1901.
Cabannes, V., Arnal, C., Bouaziz, W., Yang, A., Charton, and Kempe, J. (2024). Iteration head: mechanistic study of chain-of-thought, arXiv preprint arXiv:2406.02/18.
Chen, S. and Li, Y. (2024). Provably learning multi-head attention layer. arXiv preprint WiXiv:2402.04088
Chen, S., Sheen, H., Wang, T. and Yang. (2024). Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.1944/
Chen, S., Yang, D., Li, J., Wang. S., Yang, Z. and Wang, Z. (2022). Adaptive model design for markov decision process. In International Conference on Machine Learning. PMLR.
Chen, X. and Zou, D. (2024). What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404.01601.
Cheng, X., Chen, Y. and Sra, (2023). Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06552
Collins, I.., Parulekar, A., Mokhtari, A., Sanghavi, S. and Shakkottai, S. (2024). In-context learn- ing with transformers: Softmax attention adapts function lipschitzness. arXiv preprint arXiv:2402.11639
Deora, P., Ghaderi, R., Taheri, H. and Thrampoulidis, C. (2023). On the optimization and generaliza- tion of multi-head attention. arXiv preprint arXiv: 12.0.1280.
Edelman, B. L., Edelman, E., Goel, S., Malach, E. and Tsilivis, N. (2024). The evolution of statistical induction heads: In-context learning markov chains, arXiv preprint arXiv:2402.11006
Edelman, B. L., Goel, S., Kakade, S. and Zhang, C. (2022). Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, PMLR.
12

---

Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph. N., Mann, B., Askell, A., Bai. Y., Chen, A.. Conerly, et al. (2021). A mathematical framework for transformer circuits. Transformer Circuits Thread, 1 1.
Friedman, D., Wettig, A. and Chen, D. (2024). Learning transformer programs. Advances in Neural Information Processing Systems, 36.
Fu, D., Chen, T.-Q.. Jia, R. and Sharan, V. (2023). Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.170S:
Giannou, A., Rajput, S., Sohn, J.-Y., Lee, K., Lee, J. D. and Papailiopoulos. D. (2023). Looped trans- formers as programmable computers. In Proceedings of the 40th International Conference on Machine Learning (A. Krause, Brunskill, Cho, B. Engelhardt, S. Sabato and J. Scarlett, eds.), vol. 202 of Proceedings of Machine Learning Research. PMLR.
Giannou, A., Yang, L., Wang, T., Papailiopoulos, D. and Lee, J. D. (2024). How well can transform- emulate in-context newton's ss method? arXiv preprint 1xXiv:2433.03/83.
Guo, T., Hu, W., Mei, S., Wang. H., Xiong, C.,, Savarese, S. and Bai, Y. (2023). How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint ar arXiv:2310.
He, J., Chen, S., Zhang, F. and Yang, Z. (2024), From words to actions: Unveiling the theoretical underpinnings of llm-driven autonomous systems. arXiv preprini arXiv:2405. 19883.
He, P., Liu, X., Gao, and Chen, W. (2020). Deberta: Decoding-enhannn bert with disentangled attention. arXiv preprint arXiv:2006.0365:
Honovich, O., Shaham, U., Bowman, S. R. and Levy, O. (2022). Instruction induction: From few examples natural language task descriptions. arXiv preprint axXiv:2205.1078:
Huang, Y., Cheng, Y. and Liang, Y. (2023). In-context convergence of transformers. arXiv preprint arXiv:2310.05249.
Jelassi, S., Sander, M. and Li, Y. (2022). Vision transformers provably learn spatial structure. Ad- vances in Neural Information Processing Systems, 35 37822-37836.
Jeon, H. J.. Lee, J. D., Lei, Q. and Van Roy, B. (2024). An information -theorrtic analysis of in- context learning. arXiv preprint arXiv:2401.15550
Kim, J. and Suzuki, T. (2024). Transformers learn nonlinear features in context: Nonconvex mean- field dynamics on the attention landscape. arXiv preprint arXiv:2402.0/258
Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. and Oymak, S. (2024). Mechanics of next token pre- diction with self-attention. In International Conference on Artificial Intelligence and Statistics. PMLR.
Li, Y., Li, Y--F and Risteski, A. (2023). How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint 2xX1:23333002255
Lin, L., Bai, Y. and Mei, S. (2023). Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. arXiv preprint arXiv:2310.08566
Liu, B., Ash, J., Goel, S.. Krishnamurthy, A. and Zhang, C. (2022). Transformers learn shortcuts to automata. ArXiv, abs/2210.10049.
Mahankali, A., Hashimoto, T. B. and Ma, T. (2023). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention, arXiv preprint arXiv:2307.035: 035603376...
Makkuva, V., Bondaschi, M., Ekbote, C., Girish, A., Nagle, A., Kim, H. and Gastpar, M. (2024a). Local to global: Learning dynamics and effect of initialization for transformers. arXiv preprint arXiv:2406.038 03072.
Makkuva, A. V. Bondaschi, M., Girish, A., Nagle, A., Jaggi, M., Kim, H. and Gastpar, M. (2024b). Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint 402.:2240..0411..
13

---

Meyer, C. D. (2023). Matrix analysis and applied linear algebra. SIAM.
Muller, S., Hollmann, N., Arango, S. P.. Grabocka, J. and Hutter, F. (2021). Transformers can do bayesian inference. ArXiv, abs/2112.105111
Nichani, E., Damian, A. and Lee, J. D. (2024). How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14335.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y.. Chen, A. et al. (2022). In-context learning and induction heads. arXiv preprint arXiv:2209.11895
Polyanskiy, Y. and Wu, Y. (2024). Information Theory: From Coding to Learning. Cambridge University Press.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I. et al. (2019). Language models are unsupervised multitask learners. OpenAl blog, 9.
Rajaraman, N., Bondaschi, M., Ramchandran, K., Gastpar, M. and Makkuva, A. V. (2024a). Trans-
formers on markov data: Constant depth suffices. arXiv preprint arXiv:2407.17666
Rajaraman, Jiao, J. and Ramchandran, K. (2024b). Toward theory of tokenization in Ilms. arXiv preprint arXiv:2404.088 08335..
Sanford, C., Hsu, D. and Telgarsky, M. (2023). Representational strengths and limitations of trans- formers. arXiv preprint 2rxXi:23006028966.
Sheen, H., Chen, S., Wang, T. and Zhou, H. H. (2024). Implicit regularization of gradient flow on one-layer softmax attention. arXiv preprint arXiv:2403.08699
Sinii, V., Nikulin, A., Kurenkov, V., Zisman, I. and Kolesnikov, S. (2023). In-context reinforcement learning for variable action spaces. arXiv preprint arXiv:2312.13327.
Song, and Zhong, Y. (2023). Uncovering hidden geometry in transformers via disentangling position and context. arXiv preprint I..0422310.0866.
Tarzanagh, D. A.. Li, Y., Thrampoulidis, C. and Oymak, S. (2023a). Transformers as support vector machines. ArXiv, has/2308.16893
Tarzanagh, D. A., Li, Y., Zhang, X. and Oymak, S. (2023b). Max-margin token selection in attention mechanism. arXiv preprint arXiv:2306. 13596.
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A. et al. (2023). Gemini: a family of highly capable multimodal mod- els. arXiv preprint arXiv:231..1180.
Thrampoulidis, C. (2024). Implicit bias of next-token prediction. arXiv preprint arXiv:2402.18551
Tian, Y., Wang. Y.. Chen, B. and Du S. (2023a). Scan and snap: Understanding training dynamics and token composition in l-layer transformer. arXiv preprint arXiv:2305. 16380.
Tian, Y., Wang, Y., Zhang, Z., Chen, and Du, (2023b). Joma: Demystifying multilayer trans- formers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00533.
Vasudeva, B., Deora, P. and Thrampoulidis, C. (2024). Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.057.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. and Polosukhin, (2017). Attention is all you need. Advances in neural information processing systems, 30.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A. and Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In International Conference on Machine Learning. PMLR.
Wang, K., Variengien, A., Conmy, A., Shlegeris, B. and Steinhardt, J. (2022). Interpretability in the wild: circuit for indirect object identification in gpt-2 small. arXiv preprint 2tt.:22211.00933.
14

---

Wei, J., Bosma, M., Zhao, V. Y. Guu, K., Yu, A. W., Lester, B., Du, N., Dai, A. M. and Le, Q. V. (2021). Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652.
Wei, Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D. et al. (2022).
Chain-of-thought prompting elicits reasoning large language models. Advances in neural information processing systems, 35 24824 24837.
Wu, J., Zou, D., Chen, Z., Braverman, V., Gu, Q. and Bartlett, P. L. (2023). How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.0839)
Xie, S. M., Raghunathan, A., Liang, P. and Ma, (2021). An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.0208)
Zhang, R., Frei, S. and Bartlett, P. L. (2023a). Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.044 09927.
Zhang, Y., Liu, B., Cai, Q.. Wang, and Wang. Z. (2022). An analysis attention via the lens of exchangeability and latent variable models. arXiv preprint arXiv:2212.14852.
Zhang, Y., Zhang, F., Yang, Z. and Wang, Z. (2023b). What and how does in-context learning learn?
bayesian model averaging, parameterization, and generalization. arXiv preprini 2xXv:2305.19420
Zhou, D., Schärli, N., Hou, L.,, Wei, J., Scales, N., Wang, X., Schuurmans, D., Cui, C., Bousquet, O., Le, Q et al. (2022). Least-to-most prompting enables complex reasoning in large language models. arXiv preprint a05i::205..10655.
15

---

Contents
Introduction
2
Problem Setup: In-Context Learning of Markov Chains
2.1 In-Context Learning and n-Gram Markov Chains
2.2
Two-Layer Transformer Model
3
Theoretical Results
3.1
Generalized Induction Head Mechanism for Learning 7-- n-Gram Markov Chains
3.2
Convergence Guarantee of Gradient Flow
4
Conclusion and Future Work
10
5
Acknowledgement
11
A Related Works
17
Experiments
18
B.1 Training with Stage Splitting
19
B.2 Training without Stage Splitting
20
B.3 Prior and Length Generalization
C Additional Background and Discussions
22
C.1 How Does Transformer Implement the GIH Mechanism?
22
C,2 Feed-Forward Network for Polynomial Kernel
23
C33 Perron-Frobening Theorem
24
C,4 Sequential CE Loss
25
C.5 Standard x2 --iierrgence and Mutual Information
26
C.6 More Details on the Generalized Induction Head Mechanism
26
C.7 Further Discussions on the GIH Mechanism
26
C.8 Conclusion and Future Directions
27
Proof Sketch
27
D.I
Simplification of the Transformer Model at Initialization
28
D.2 Analysis for Training the FFN and the First Attention Layer
28
D.2.1 Training the FFN: Identification of the Information Set S*
29
D.2.2 Training the First Attention Layer: Convergence (w)(h) to One-Hot Vector
30
D.3 Analysis for the Training of the Second Attention Layer
31
Analysis of the Training Dyanamics
32
E.1 Conditions on the Sequence Length
32
E.2 Analysis for Stage
33
E.2.1 Additional Proofs for the Stage
37
16

---

E.3
Analysis for Stage II
38
E.3.1
Additional Proofs for Stage II
45
E,4
Analysis for Stage III
45
E.4.1 Additional Proofs for ................................
53
E,5
Lemma on GIH Approximation Error
55
F
Auxiliary Lemmas and Their Proofs
57
Fl1
Useful Inequalities
57
F.2 Approximation Errors for Dynamies Analysis
58
F3 Lemmas on Concentration of Markov Chain
73
Organization of The Appendix
The appendices are organized as follows:
In A, we present in-depth discussion the related works.
In $B, we discuss the experimental details.
In &C, we discuss the implementation of GIH mechanism, provide explicit expressions for the FFN realizing a low-degree polynomial kernel, and review basics related to concepts mentioned the main text.
In $D, we provide a high-level overview of the proof of our main results.
In SE, we present the proof for Theorem 3.6.
In $F, collect auxiliary results used in the proof of Theorem 3.6.
A Related Works
In Context Learning (ICL). Commercial Large Language Models (LLMs) such as ChatGPT (Brown et al., 2020), GPT-4 (Achiam et al., 2023), and Gemini (Team et al., 2023) typically operate an autoregressive manner. These models exhibit remarkable ICL capabilities, without requiring further training. Previous research explores various aspects of the in-context learning (ICL) ability of these models. This includes their performance in zero-shot and few-shot learning scenarios (Honovich et al., 2022; Wei et al., 2021), the use of the chain of thought method to enhance reasoning (Wei et al., 2022; Zhou et al., 2022), and learning with multi-modalin research highlights the properties and advantages of using transformers beyond the traditional ICL setting, thereby broadening our understanding of their capabilities and applications (Edelman et al., 2022: Li et al., 2023; Jelassi al., 2022: Sanford al., 2023; Giannou al., 2023; Liu al., 2022: Tarzanagh et al., 2023a,b; Tian et al., 2023b,a; Song and Zhong, 2023; Deora et al., 2023; Chen and Li, 2024; Rajaraman et al., 2024b).
(Alayrac et al., 2022). Moreover, recent
There a large and growing body of literature on understanding how transformer architecture enables ICL. One strand of research proposes to understand ICL by casting it as a version of Bayesian inference expressed by the transformer architecture. See, e.g., Xie al. (2021); Muller al. (2021); Zhang et al. (2022, 2023b); Ahuja et al. (2023); Jeon et al. (2024); He al. (2024) and the references therein. Another line of work investigates how transformers internally emulate specific algorithms to solve ICL tasks, where Akyürek al. (2023); Von Oswald al. (2023); Fu et (2023); Ahn et al. (2023); Mahankali et al. (2023); Giannou et al. (2024); Wu et al. (2023) focus on learning with linear regression tasks and Bai et al. (2023); Cheng al. (2023); Collins et al. (2024); Guo al. (2023) investigate transformers' capabilities in learning with nonlinear functions. However, all of these works above focus regression tasks where token (or token pairs) in the prompt sequences are i.i.d. or uncorrelated, which may not capture the more sophisticated data structures in real- world applications.
In addition, to study ICL. with correlated data, there is also substantial interest in understanding how ICL operates over data drawn from Markov chains, providing insight into how transformer architectures contribute to ICL in these settings (Edelman al., 2024; Makkuva et al., 2024b; Chen
17

---

and Zou, 2024). Furthermore, Lin et al. (2023); Sinii et al. (2023) show how transformers can solve reinforcement learning problems in an in-context fashion.
While many of the aforementioned works focus on the expressivity of the transformer model on dif- ferent ICL. tasks and the statistical properties of the learned models, understanding training dynamics from an optimization perspective is also crucial for comprehending ICL by transformers. The training dynamics for one-layer attention models have been investi gated under different data models for both regression and classification tasks (Zhang et al., 2023a; Huang al., 2023; Tarzanagh et al., 2023a,b; Kim and Suzuki, 2024; Chen et al., 2024; Vasudeva et al., 2024; Li et al., 2024; Thrampoulidis, 2024; Sheen et al., 2024). These studies offer a thorough characterization of the training process, yet they have limitations - they are not directly applicable to data drawn from Markov processes and are confined to single-layer attention. Our work belongs to this line of research and we adopt a awo-attention-layer transformer architecture, which is more complicated than the transformer studied in these works.
Induction Head. Elhage et al. (2021) introduce the concept of "induction heads" as the mechanism underlying the ICL capabilities of transformers. Since then, there has been surge of interest in understanding the induction head mechanism and its role ICL. At a high level, the induction head mechanism works by matching the history of the current token with those seen previously in the sequence and then predicting the next token based on the matched historical sub-sequences. Olsson et al. (2022) provide empirical evidence highlighting that induction heads are crucial in facilitating the ICL capabilities of transformers. Bietti et al. (2024); Edelman et al. (2024) conduct further empirical investigation into the development of induction heads specifically tailored for the ICL of bi-gram data models. Rajaraman et al. (2024a) provide explicit constructions of single-head transformers with constant depths that can learn n--ram data. Also, a wider range of functionalities exhibited by induction heads that interact with various other mechanisms have been observed by Wang et al. (2022).
From a theoretical perspective, Nichani et al. (2024) study the ICL of first-order Markov chains using a two-layer transformer and demonstrate the formation of the induction head mechanism. Makkuva al. (2024a) also prove that training a single layer attention with feed-forward layer on first-order Markov data (with {0. 1} vocabulary) can converge to either to global or local minima depending on the initialization. However, the first-order assumption seems to be quite restrictive, especially when modeling the natural language, where the tokens can depend on multiple previous tokens. Most related to our work is Nichani et al. (2024), where they analyzed how training by gradient descent enables two-layer transformer to learn the latent causal graph underlying the ICL data. However, the analysis in Nichani et al. (2024) applies to Markov chains where each token has at most one parent, and it remains unclear how extend the analysis to more general n-gram Markov chains.
In this work, we show that generalized version of the induction head mechanism can emerge when training two-layer transformer on n-gram Markov chains. Moreover, our transformer models are more sophisticated, incorporating features like relative positional embedding, multi-head atten- tion, an FNN layer, and normalization. Notably, we provide an in-depth dynamics analysis of the corresponding FFN layer and two-layer multi-head attention.
B Experiments
In this section, we first detail the setup for the experiment in Figure 3, and then provide additional results for training a model that also incorporates the word embedding matrices Wo, WK. Wy and the output embedding matrix Wo in the first attention layer. Let first detail the data setup that is used for all the experiments in this work.
Data generation. The dataset for the ICL task is generated as n-gram Markov chains as described in $2.1. We take pa = (-1, -2} the parent set. Thus, the number of parents is = 2 and the token embedding dimension d = Note that for each sequence, the transition matrix π <<x Upa) is of shape d x d". We assign prior distribution P for the transition matrix, which is defined such that each column of the transition matrix of kernel π is independently drawn from a symmetric Dirichlet distribution with parameter a = 0.01, i.e., T- ~ Dir(a 1a). Note that each chain has different transition kernel but follows the same prior distribution P. We randomly sample 10,000 Markov chains with L = 100 from the prior distribution P; 9,000 are used for training and 1,000 for validation.
18

---

Stage E: Train
Siage Train wP)
Stage III: Train
1.0
0.8
0.4
0.0
Epoch
Beginming
Ending
Epoch
(c)
Figure 4: An illustration of the transformer parameters during the three-stage training. This is the same figure as Figure We train a transformer in TF(M = 3, H 3, d 3, D = 2) with I, = 100, pa (-1.-2}. In (a) we show the evolution of SH))< in the first stage of binary coding in {0, 1)3 to indicate each subset S. Recall that "110" represents = {1, 2). which is exactly S This figure shows that Ps* gradually increases to one while the any other ps decays to zero. In (b) we plot the RPE weights of the first attention layer before and after the second stage of training. Here the h-th column corresponds to the RPE weight vector of head h. This figure shows that w and w increase to large number after training, while w13) stays close to its initial value. Thus, we have a (w)()) = (() = 1. That is, the first two heads are trained to attend to parents -1 and -2, respectively. In (c) we plot the evolution
training where ps c D .g.. We
in the last stage of training. This figure clearly exhibits a two-step growth pattern and keeps increasing throughout this stage. In summary, the results of the simulation experiments coincide with the theoretical results.
B.1 Training with Stage Splitting
present the simulation results with model TF(M, H, D) in (2.5) and training in the three-stage manner. configure the model with window size M 3, number of heads 3, vocabulary size d = 3 and maximal FFN degree D = 2.
Model initialization. The RPE weight matrix w is initialized such that the (-i)-th diagonal of !!lh) has value th for i 1,2,..... M, while all other entries initialized -o. See Figure for interpretation. We initialize (h and set the remaining entries within the size-M window to 0.01 to ensure symmetrization breaking and some initial correspondence between heads and
parents. For the FFN layer that learns the polynomial features, es for S E [H]&D are initialized to 0.01. The initial value of a in the second attention layer is set to 0.01.
Training settings. The models are trained using gradient descent with respect to the cross-entropy
loss and constant learning rate that is set to one for all stages. We train the model in Stage (update parameters {cs} only) for 2000 epochs, in Stage II (update parameters (w)) only) for 50,000 epochs, and in Stage III (update parameter only) for 5000 epochs, respectively. All experiments conducted using a single Nvidia A100 GPU. The results are shown in Figure 4, which matches our theoretical results.
19

---

0.000100
train
1.0
0.020095
000
001
0.000090
011
0.000085
0.4
101
0.000080
110
0.2
0.000075
4000
10000
2000
1000
6000
8000
10000
Epoch
Epoch
(a)
(b)
1.0
2
0.8
0.8
0.5
0.6
0.6
0.3
0.4
0.4
0.2
0.2
0.2
8
0.0
2000
4000
6000
0000
10000
Beginning
Ending
Epoch
(c)
(d)
An illustration of the evolution of gradient descent dynamics when training trans- former model specified in $B.2 with word embedding matrices {Wo, Wr, W. Here the dynamics are not split into three stages and each gradient descent step updates all parameters. We set . = H 3. d = and D 2, the number of input token is L 100, and Markov chain has parent set pa = (-1, 22.. In (a) we show the training loss of the model, which shows that the loss decreases and converges to some value. In (b) we show the evolution of ps where we use binary coding {0, 1]3 indicate each subset S. Here, ps+ has code "110", which corresponds to the true parent set. This figure shows that initially a wrong Ps dominates at the early stage of training, which corresponds Ps- increases and becomes dominant. However, ps- does not increase to one and is about 0.6, and there two ps that are about 0.2. In (c) we show the RPE weights of the first attention layer before and after training. The entries corresponding to the true parents, w and (2)
Figure 5:
= {2,3) (code '"')"". Then eventually
significantly increase after training, while w w slightly increases from initialization. This figure shows that each attention head focuses on copying a single previous token. In (d) we show the evolution of the weight in the second attention layer. We observe similar "elbow" curve in Figure 3-(C).
(3)
B.2 Training without Stage Splitting
Previously in $B.1, we show the simulation results on the simplified model (2.5). Now we present the results of additional experiments based on the full model defined as follows.
First Attention:
(M) = o (X ((()) XT + w xw
E [((+))xd,
Concatenate & Normalize:
V = LN((D ,...,,......,.,, (,,),
E [(L+1)x(H++))
FFN & Normalize:
Ò d(V)/vo)
E ([++))xde
Concatenate
X' = x'=
E ((L+))(((
Second Attention:
Y = aaaa (Tt+1) ((x),)))
E (l++1)xd
In head the first attention layer, w-(h) is the relative positional embedding matrix, and we include w-(h) (h) dddd w-Oh) (h) gdxd and w(( e ddddd as the weight matrices for the query, key, and value
20

---

a pre-trained transformer with = = 3, = 3, and 2. These are the parameters in of the first attention head in the first attention layer. Since d all word embedding matrices shape As shown in (a) and (b), w and ( do not change much compared to their initialization value 0.001. Thus, they are both close to the zero matrix and play negligible projections, respectively. That is, in the full model, we the attention heads has more weight matrices than the simplified model. Another difference is that we also explicitly include the residual link that copies X to the output of the first attention layer. For the FFN layer, Φ RR - Rd is the same feed-forward network specified in (2.3). Here, we use standard ey-layer-normalizatinn LN(-), defined as
LN((2, v))
1/
1
The second attention layer takes as the value, which comes from the residual link (i.e., concatenation of U and V while X in remains the same after (2-normalization) In comparison to the simplified model in (2.5), here incorporate the query, key and value projections for the first layer as in standard transformer architecture.
Our training setup is similar to that in $B.1. We use the same dataset and a similar training settings.
All these weight matrices (-(h) wrh)) and w(()) are initialized as identity matrices scaled by 0.001.
Q
We initialized the RPE vector whth as w
for / 1, 2, 3, and leave the remaining entries within the length-M window to 0.01. We trained the model with all parameters together for 10,000 epochs with the same loss function and learning rate. As illustrated in Figure the full model converged to state comparable to our simplified model. We further plot the w w(( w for the first head after Q training in Figure The results demonstrate that the model converges to a point where the query and key projections are close to zero, which leaves the RPE weights to dominate the attention mechanism.
This fact justifies our simplification in (2.5) where we remove the query and key projection weights and set w-(h) to be identity matrix.
4
W
W
0.0100
-0.5
0.2100
-0.5
0.0
0.0
0.0075
0.0
0.06
0.0050
0.0050
0.5
0.0025
0.5
10-
0.0000
0.02
15
1.5
-0 00.0025
1.5
-0.0050
-0.0050
2.0
0.0075
2.0
-0.0075
-0.02
2.5
0.00000
2.5
2
0
(a)
(b)
Figure 6: visualization the word embedding matrices W(1), W W
of
role in the first attention layer. Besides, in (c) we plot W(Y) which establishes a clear diagonal structure, with the diagonal entries growing to 0.07 compared to the initialization value 0.001. Thus, ((1)) proportional to the identity matrix.
B.3 Prior and Length Generalization
We further test the model learned by the three-stage strge-stage training on sequences coming from different priors and of different lengths. Note that our pre-trained transformer learns to perform GIH. As introduced in 13.1. 1, the GIH estimator be applied sequence with arbitrary length and does not concern the prior distribution the underlying Markov chain. Thus, it is natural to see the pre-trained transformer can also generalize to different lengths and prior distributions.
Recall that train the transformer model with sequence length L. 100 and the concentration parameter of the Dirichlet prior is a = 0.01. Here, we test the pre-trained transformer on new sequences of different lengths and sampled from different prior distributions. That is, with a different
21

---

concentration parameter sample a random Markov chain, and generate sequence of length L, and evaluate cross-entropy loss for predicting "L+1. Here choose a E {0.05, 0.1, 0.2) and range L. from 10 to 1000. When generating the data, the Markov chains share the same parent set pa {-1, -2 with the pre-training data. The results shown in Figure 7. The results show decreasing trend in testing loss as the sequence length increases. For a = 0.2, we observe first a small increase in the test loss when I just exceeds 100, but then the loss decreases as I. increases further. This experiment shows that the pre-trained transformer indeed generalizes in length and robust to the change of prior distribution.
le-4
d=0.1
Volidetion
7.0
6.5
6.0
10a
10a
10a
Sequence length
Figure 7: Generalization capability of our model different sequence lengths and prior distributions. We plot the cross-entropy loss of the pre-trained transformer model on se- quences with different lengths sampled from Markov chains with different prior distribu- tions. The prior is Dirichlet distribution with a E {0.05, 0.2} and we vary the length L in {10, 20, 50, 100, 200, 400, 700, 1000). The pre-training data contains sequences of length L 100 and a 0.01. For different a, we see that the error has decreasing trend L increases. This shows that the pre-trained transformer can generalize in length and is robust to the distributional shift due change of prior.
C Additional Background and Discussions
C.1 How Does Transformer Implement the GIH Mechanism?
In the following, we briefly illustrate how a two-attention- layer transformer model as introduced in (2.5) implements the GIH mechanism. As we will show in $3.2, gradient flow with respect to the cross-entropy loss converges this transformer in the limit.
X
X1
Value
L+1
FFN & LN 8 = (1,2)
U1
Key
Query
Head
X1-3
Residual
Head
X1-2
X1
Head
Ul
u U
X/-1
X1-3
X1-2
X
X1
X2
0
Figure 8: Illustration of the GIH mechanism in two-attention-layer transformer model. Here, pa = (-1,-2) M = and S* = {1, 2). The first attention layer copies the parents (including the information set S*) the current position. Then the FFN layer together with layer normaliza- tion generates the features u, using the parent tokens within the information set S* The second attention layer treats each I the value, and aggregates keys and query that come from the learned features using the attention mechanism. The + 1-th token is padded with zeros in the input.
the prediction by matching the
Step I: The First Attention Layer Copies the Information Set S* to the Current Position. Suppose the number of heads is equal to the window size for simplicity, i.e., H = M. Then, attention
22

---

head h E S can attend to the h-th parent token by setting the RPE weights in the softmax function
to be (h)) = p e-h for sufficiently large p. where e-h E IR is the canonical basis vector with the (M 1 h- entry being one and all other entries being zero, As result, each t) (h)
(h) for h S*
satisfies 2 = h1
Step II: FFN Generates the Polynomial Features of the Information Set S As have intro- duced in (2.3), each learnable es in the FFN layer determines the contribution of the corresponding subset S to the output feature. To let the optimal information set S dominate the output, we set cs. whereas cs for S + St. The exact form of the output the FFN layer, ((mr), deferred to $C.2. Here the only property we require is that
st (6(0)) ((LL)))
II (o) (1) (11
1 XL+1-S.).
(C.1)
AES
Here X1-s- == ( 8 S*) and in the last equation we use the orthogonality and normalization of the vocabulary embeddings.
Step III:
The Second Attention Layer Aggregates Tokens with Matching History on S We can interpret é in (C.1) as an indicator for whether the information set of a token H; matches the information set of the token F1+11 Then for the second attention layer, by setting be sufficiently large, the output will become
exp(a 81) -31
N-1. AN--EEmae Ei=M41 Tt 1(X--s- = -L+1-S+), 21, (L- Σi= v1, otherwise,
1=M+1 Ek-m+1 exp(a. st)
(C.2)
where N Et- (Xt-s- -r.+ That is, if least one token H; has matching information set 11+1, i.e., their histories restricted to S the same, the second attention layer outputs the average of such tokens. Otherwise, it outputs the average of previous tokens from EMS+1 to Iz. In Lemma E.6 in the appendix, will show that the model learned by gradient flow implements the GIH mechanism up to diminishing approximation error.
The weights of the transformer constructed above are illustrated in Figure 9. We consider the trans- former model with M H = 3, d = 3, and D 2. In this case, in the first attention layer, for each h E (3). ( has three finite parameters w w and w By our construction, we have w p for all [3] and the rest of the entries of (w } are all equal to zero. In Figure 9-(a)
plot the top ten by ten block of w where _ p is shown in yellow and w w w_3 are shown in purple The gray color stands for -oo entries. Figure 9-(b) we plot (w(h) jaj In Figure 9-(c) we plot the parameters of the FFN. Since H 3 and D 2, [H}D contains seven elements: Ø. {1), (2), {3). (1, 2...... (1, and {2, 3). We use binary strings of length 3 to index these seven subsets, where the i-th bit indicates whether element i is included in the subset. For instance, " "110" represents {1,2]. We set S = (1,2), cs. 1, and cs = for any other S.
(1)
C.2 Ford-Forward Network for Polynomial Kernel
Lemma C.1. Recall the FFN satisfying (2.3), which maps a vector 2 E RdHH to a vector in Rdd, We write as (2 ...., .,,2(1) (T) where (h)) Rd for all h [H]. Let (h) be the i-th entry of (h), Then we can explicitly construct ((.) by letting
((2 (0..,,.11)) (1)) ((E,))...,,11)))
cs II hES
(h) 2/ finhees S [d Se { {0]
(C.3)
which is equivalent to
& ((2 (..,, (B1)) (cs vee (ec( SE(M)gD
where veol.) is the vectorization operator that transforms a tensor into vector by stacking all the entries in the tensor. That is, for any S, consider the (S] vectors in RRd, (2)))
Thes. In (C.3)
23

---

kx(3)
W
1.0
Gray represents 00
0.8
Cgt dominates
3
2
0.3
000
001
010
011
100
101
110
(4)
(b)
(c)
Figure 9: Limiting model of TF(M = 3, 3. = 3, = 2) that implements the GIH mechanism with L 100. pa (-1,-2) (a): The top left 10 by 10 block of w that attends to the -1 parent. (b): The RPE weight heatmap for all 3 heads, where the h-th column corresponds to the RPE weight vector of head h. (c): In the GIH mechanism, only one c' for the optimal information {0, 1 8 to indicate each subset S. Here, S = {1. 2} is the parent set, which is represented by " "110".
dominates. For the label of the x-axis, we use binary coding
compute all possible products of the entries of these vectors and multiply them by cs. In particular, for each S E [H) we enumerate in E [d] for all h e S. Therefore, the output dimension of Φ is given by
d. Σ
di
(C.4)
SE(H
Proof. First, we that the indices of o(-) have grouped structure first enumerate all subsets in [D] and then enumerate all monomials with superscripts in S. Since there are dis monomials, the output dimension given by (C.4).
It remains to verify (2.3) with o(-) defined in (C.3). To this end, we note that for any u, E RdH and any S E [H] L.D. we have
Σ
II
(1)
(N)
II
Σ
(h)
(b)
(w ((h)).
U.
Ui
UiA
=
(dd)esE[dis))
hES
hES
heS
which directly implies (2.3). Therefore, we conclude the proof of this lemma.
C.3 Perron-Frobenius Theorem
Next, we review the basics for the celebrated (Meyer, 2023, Chapter 7). We consider the following class of irreducible matrices.
theorem on non-negative matrices
Definition C.2 (Irreducible Matrix). A non-negative square matrix P Idxd is called irreducible the induced directed graph G is strongly connected, ie., for any pair of nodes the graph, there always exists directed path that connects these two nodes. Here, the induced graph G is defined based on d nodes with adjacent matrix A given by Aij 11 (133 # O).
In particular, if is a stochastic matrix that corresponds to d-state Markov chain, then starting from any state, we reach any other state with positive probability in finite number of steps. The irreducibility property also has an equivalent definition in the matrix form. That is, any permutation matrix T. TPT-1 cannot be written as an upper triangular block matrix with the following form
Mi
M2 Ms
In other words, an irreducible matrix does not have a nontrivial absorbing subspace that aligns with the standard basis.
24

---

In this work, we require more than the irreducibility property from the transition matrix Pr defined in $3.2. In fact, we need the existence of unique stationary distribution (which is not guaranteed by the irreducibility) so that the chain has a sufficiently fast mixing rate. This enables us to learn with a finite sequence length L. To achieve this, one typically needs the second largest magnitude of the cigenvalues of Pa. denoted by X, to be bounded away from 1, which is the leading eigenvalue of the transition matrix. The difference - À is also referred to as the spectral gap. It is all the entries of P= are positive, then P7 irreducible and there is only one leading eigenvalue on the spectral circle with the corresponding cigenvector given by the chain's 's stationary distribution μ²" and all the other eigenvalues have magnitude strictly less than 1. However, for our case, the transition matrix Ps has zero entries by definition. Fortunately, the nice property on the existence of spectral gap can be generalized to class called primitive matrix.
that if
Definition C.3 (Primitive Matrix), A nonnegative and irreducible square matrix P iss called primitive if there exists an integer k such that all the entries of pk are positive.
By definition of the primitive matrix, one can immediately see that for any k 4 k, Pkl is a positive matrix. The following is the celebrated Perron-Frobenius theorem that characterizes the spectral structure of the primitive matrices.
Theorem C.4 the following stalements hold:
Theorem for Primitive Matrices). Let P be primitive matrix. Then
1. The leading eigenvalue of P is real and positive, and it is the unique eigenvalue with the largest magnitude. In particular, if P is stochastic matrix, then the leading eigenvalue is 1.
2. The leading eigenvector of P is positive and unique up to a scaling factor. In particular, if P a stochastic matrix, then the leading eigenvector is the stationary distribution of the Markov chain with transition kernel P.
The
theorem guarantees the existence of a unique stationary distribution μ² when the transition matrix Px is primitive. In particular, when we further assume that the transition matrix Px has a spectral gap, the chain is sufficiently mixed, meaning that we can thus approximate sum over the entire sequence with average with respect to the stationary distribution. In particular, the approximation error will decays with the sequence length L.
C.4 Sequential CE Loss
In this work, we only consider the prediction error on the last token in the sequence as in (2.1):
[(fer) = -..-EEPP... (1 [log (fr((14411 (II...) + +-
In practice however, people often train the transformer model by minimizing the cross-entropy (CE) loss over the entire sequence. We demonstrate that our analysis can be extended to training on the entire sequence. In this vein, we define the sequential CE loss as
= Σ -EASP, [log (foe(fer 214) + +)].
Leeg (frr)
(C.5)
1=1
One can equivalently view this sequential CE loss as an aggregation of the CE loss for sequence length ranging from 1 to L. We argue from the following two perspectives that our analysis can be extended to the sequential cross-entropy (CE) loss:
1. Due the use of relative positional embedding (RPE), the transformer's predictions are invariant to the absolute positions of tokens within a sequence. Intuitively, this implies that even if we choose different sequence length L', the model can still handle the task in the same manner.
2. By Assumption 3.5. the chain is sufficiently mixed for large L. In the analysis, we actually use XI-Mal
M.l. = (att, It-1, , ~ 10²,, where μzx is the stationary distribution over a Length-(M+1) window, to approximate the aggregation over Xi- M- for 1 = M 1,,..,, ...... L
in the sequence. For example. this approximation is reflected in the transition from (D.3) to (D.4) in the proof sketch in $D. Since changing the sequence length does not affect the underlying stationary distribution, the only issue the approximation error. In particular, for sufficiently large L, the CE loss at large constitutes the majority of the sequential CE loss in (C.55, making the CE loss at small | negligible.
25

---

C.5 Standard x² -Divergence and Mutual Information
The divdivegeence (or disttancce between probability distributions P and in the same probability space is defined as:
(P(x) - -(x)) ((x)
Dx (PI()) = Σ
zEsup((Q)
where the summation is taken over all elements x the sample space where Q(x) > O. The x².
mutual information between two random variables X and Y with joint distribution Pxy and marginal distributions Px and Py is defined as:
(x; = Px=(Pxxxll P Pp) 1
where Px Py is the product of the marginals, meaning (Px & Pr) Pr())))(y)) For Markov chain + Z, the x2 -mutual information satisfies the data processing inequality
[x=(X; S (Y=
which follows from the observation that x² -dvieergee is also an f-divergence.
C.6 More Details the Generalized Induction Head Mechanism
Recall that we define the Generalized Induction Head (GIH) estimator in (3.3). Specifically, GIH(TIL1; M, D) is constructed in two steps. First, we find the information-optimal subset S of [M] by solving (3.2). Second, we build d-class kernel classifier predict IL+1, where the "data" used by such classifier are {1bs- (4), 21 he [1.++ Here { (I), E [M +1, + 1}} are features constructed at each position based on the partial history given S In particular, similar to (C.3), for any subset of [M]. any input token sequence # and any position / [M + 1, + 1]. we define (I) = t (1; L as
ss(()= vec 2 -1-.) - -
(I (m)-s); -
tistes S (ld) E R
#ES
In other word, sls is given by expanding the rank-1 tensor spanned by (xi-*) into vector. Here Zi-s E X is vector in Rd and we let (zı-s);, denote its i,-th entry. The rationale behind (I) similar to ø introduced in (C.3). We form long vector containing all the products of the entries of vectors { Here we omit the dependency of ths on the input sequence T1:L to simplify the notation. Furthermore, ths induces polynomial kernel such that for any 1, m E [M L + 1). have
(Vs(1), (s(m) = II (.) Tm 11 1[ Im-s, Vs E S}.
#ES
That is, feature ths selects the token position pairs (1, m) such that the partial histories induced by S at position and m are exactly the same.
Based on si (I), zi heeMMM4,,)) GIH forms kernel classifier using the indicator kernel. Specifically, for any j E [d]. by (3.3), GIH (1 M, D) outputs each ej E A' with probability
(P D, = =e;) Etam 1(5i- -1141 Vs E S*) 1 = Eÿ}
Eme 1 [ L+1- Vo E S.
C.7 Further Discussions on the GIH Mechanism
We conclude this section with further discussions on the modified x" mmmuul information and low- degree polynomial kernel for the FFN within the GIH mechanism.
On the Modified x 2-Mutlll Information. Now that we have shown how gradient flow approaches the desired GIH model, it is natural to ask the following questions: What is the optimal subset S* that the model selects? How well does the model perform? For the purpose of illustration, let us consider symmetric case where the stationary distribution μº over length-r, window is uniform
26

---

over X One can verify that in this case, the stationary distribution over window of any other length is uniform as well, and the modified mutual information can be simplified into
log Ix=(5) = log [x=(S) log c,
(C.6)
where 1x2 (S) is the standard x² mutual information between (= Z-s) and (z), and the second term log serves penalty on the model complexity. Thus, the GIH mechanism is reaching a balance between the model complexity and the information richness. Below we characterize two scenarios where the model will select the exact parent set, i.e., S* =pa.
1. If 1, i.e., each token only has one parent, then S* - pa. This is because S simultancously maximizes both terms in (C.6), thus reproducing the results in Nichani et al. (2024).
2. If is known priori and restricting the polynomial kernel to S E [H]
{s [H] |S| =
n for the FFN layer, then S pa. Here, the penalty term does not influence the selection and the exact parent set maximizes the mutual information by the data data-procsssing inequality.
In the general case, however, the model could be much more flexible, and it is possible that the model selects only a subset of the true parent set or even some non- non-parent tokens that are also informative. The rationale is that with more complex model, e.g., selecting a large S, the model are able to make more accurate predictions for large I. but may endure a large estimation error for small L, as the exact matching X = XL+1-S may appear rarely in the sequence.
On the Low-Degree Polynomial Kernel. The goal of using a low-degree polynomial kernel in (2.3) is to strike a balance between model complexity (which is also related to computational cost) and the model's s accuracy. In this regard, we have the following corollary.
Corollary C.S. We always have |S*| S n regardless of the choice of D, where S* = argmax M1 log (S) for (S) in (C.6)
The reasoning behind this corollary is follows. Consider any set S with |8| n, we have (S)(S) S Ix= (pa) as the true parent set is the most informative. Moreover, since |pal = n < |S] S suffers from a larger penalty, As a result, we have log (S) < log I== (pa) when S has more than n elements. other words, is without loss of generality to set D S n.
C.8 Conclusion and Future Directions
In this paper, we have studied the training dynamics of a two-attention-layer transformer model for learning n-gram Markov chains in an in-context way. Our theoretical analysis underscores a congruous interplay between the multihead attention mechanism, the feed-forward network, and layer normalization that yields a generalized version of the induction head mechanism during the training. In partieular, we prove that the generalized induction head mechanism adopts a modified 2m-mut information criterion for parent selection that strikes a balance between information richness and model complexity. To our best knowledge, our work gives the first theoretical evidence for learning an induction head mechanism with n-gram Markov data, which potentially sheds light on the inner workings of aarge-scale transformer models.
Our work opens new directions for developing rigorous understanding of the transformer models. natural direction would be that one can find such mechanism with standard FFN layer using multi-layer perceptron and standard layer normalization in the more practical transformer model. The intuition that our FFN layer in (2.3), which is further instantiated in (C.3), lies in the space of low-degree polynomials and can be well represented by a MIP with sufficient dimensions and proper activation functions. Initial attempts to learn nonlinear features have also been made by Kim and Suzuki (2024). Another direction is to investigate the training dynamics beyond a single loop of this induction head mechanism, e.g., iteration head with recursively refined predictions (Cabannes et al., 2024), and how the induction head mechanism oceurs in multi-layer transformer models.
D Proof Sketch
In this section, we discuss the main ingredients of analysis of gradient flow. First, we show in $D.1 how simplify the model based on choice of the initialization and the structure of the disentangled transformer. We then proceed to present the main proof idcas for the three stages of the gradient flow dynamics, where the training yields the following behaviors:
27

---

Stage I: A unique S* E { stands out such that the associated parameter cs+ dominates those of the other sets. As a result, Pi(()) = c (t)/Cn(I) approaches to one.
Stage II: For each h S*, o (w)(1) approaches one-hot vector eM-1-A RM where w(() contains the parameters RPE of the h-th head, During this stage, each head concentrates on copying a particular parent.
Stage III: Finally, a grows and reaches Oloo L). As a result, the trained model approxi- mately implements the GIH mechanism (M,
D.1 Simplification of the Transformer Model at Initialization
We first simplify the expression of the transformer model at initialization under Assumption 3.3, by showing that the attention scores of the second attention layer admit simpler form.
For the second attention layer, write the output as y o(ss) where 1 the row vector of the similarity scores. Recall from (2.5) that the FFN layer with normalization
M UIL))
outputs U (VV///CD 1xx++))dd, and we denote the /-th row of U by u e(vr)/ For 1 = + 1,..,, L, the 1-th entry of is given by
(n, #1+1) = (o(v)), (VLL+))/CD
and the other entries are all -c. By the property of the FFN layer in (2.3) and the definition CD =Sse[Hlg cs. we rewrite the above attention score as
EseMHie Cs (v (it))
for M
L.
(D.1)
ESE(H]go cs
Note that under Assumption 3.3, by the definition of Aur in (3.6), we have sufficiently large gap (h) wh for all j + at initialization. Thus, exp(w")) > exp(w!")) for all j + h, which implies the following approximation:
M exp(wi1h) (h))
Σ M
21- 2o 81 for M + +1....,I
Eラ= exppex (T))
k=1
This further implies that for = M .......... L, we have
(o) RES
(b)
(N) ULL
(m1-n- 1(21-1 1101-=2L-1-2 ZL-1-1 for S), hES
II
(D.2)
which is binary value indicating whether the query and the key token's history match the subset S. Combining (D.1) and (D.2), we obtain the following simplified expression for st:
EselH]e os 1f21= = #L+1-1 for é SE
31
EseMpED 2
Σ SE
ps 1 F1+1-; for 8}
where we denote ps es/ EseMMEp es for S [H]
| {D
In summary, when Aw sufficienlly large, w/th approximately copies the token F/-h- As a result, the attention score 81 satisfies
Σ ps 1(x TL+1-; for S}.
St
SElHIg
D.2 Analysis for Training the FFN and the First Attention Layer
The first two training stages involve the dynamics of the weights of the FFN, {cs (cs)sE(M])D> and the weights of the first attention layer, (H) ) The analyses of these two stages have similar structures and contain the following essential steps:
1. Derive the explicit expression of the dynamics of the weights, via direct calculations.
28

---

2. Unveil the key quantities (related to the modified -MI that dominantly drive the dynamics, by replacing the empirical average over the context sequence with the expectation over the stationary distribution, along with other approximations.
3. Then based on the above characterization of the dynamics, we can show the convergence of the weights to the desired values.
D.2.1 Training the FFN: Identification of the Information Set S*
In the first stage, we track the dynamics of (t) for each S E [Hgo For convenience, we drop the dependence on é in the sequel.
Recall the output of the model is y = (a(a x) and the cross-entropy loss function is [(O) Em=P,a [())), where (()) can be written as !(@) (#I,+1, log(y +1)). We ignore the small constant E in the following proof sketch for simplicity. We also abbreviate the vector of attention probabilities in the second attention layer as R
Calculation of the Dynamics of c..
By a direct calculation for the loss l and st in (D.1),
2cs Innes (u)") .it1)
(K)
0e -a - 0ila.50 8)- 8s:
FL++1
ast
20s81
(x1 - v)
dcs
Cn
Cn
Here the vector FL+1/Y is obtained by element-wis division and (a is the I-th entry of o s). Then applying the chain rule, we obtain the following dynamics for eg along the gradient flow:
O, log c =
Σ
L
ae as
4a
Σ
L
E o((a-s)
5)
II (N) (
FL+1
(xi -
Cs
Us dcs
CD
y
1= M+11
1=M+1
hES
4a
Σ
E oila 8). .1
x1+l 81
(al
v)
CD
y
L=M+1
f(t)
Note that here the second term f(1) is independent of S, and it will be canceled out when we consider the difference of the derivatives, 0 log c² - 0g log ez., for two sets S, E [H] This is why we focus on the time derivative of log dz..
Relate the Dynamics to the Modified MM by Approximations. Now using the approximation in (D.2) for [ne (h) -121)) expanding (IL+1/y) T()(() - y) coordinate wise, and noting that
(a(a s) 2e 1/(L- - -M) as we have small a in the second attention layer, we arrive at
O, log ~ (L M)C 1=M+1
4a
Σ E
1(X-s= -X+1-s)
(EL+1 = 21 ER)
1) f(t).
y(k)
(D.3)
where y(k) denotes the k-th entry of y and Xi-s x- (xn i E S) denotes the history of xq on the set S, similar for XL+1-S. Note that y(k) x (L- -)--M)--- ---- Ei=M++ 1(x; = ex) Ro 1² (ex), which follows from the mixing assumption of the Markov chain that allows to replace the average over I 1= M-1...... by by the expectation over the stationary distribution. Also for the same reason, we can replace (Ii, Xi-s), (1 (I X with two independent copies from the stationary distribution u=, i.e.,
4a
1(====eN (*
- -1 -1)]-50 f(t).
o, log ~ (,,,)) xpo = CD
(D.4)
See the approximation from 9 to Is,s in SE.2. Indeed, the first term in (D.4) becomes the modified 2-I (S), which is defined in Definition 3.1. This gives rise the following approximation:
4a - - f(t)
a log L3
(D.S)
Cp
29

---

Since the value of (t) is independent of the specific choice of set S, it is clear that the set S achieving the fastest growth rate is the information-optimal set S* srgmax Mlso (x)(S) that maximizes the modified x- 2MI within {D
Convergence of pg.. Note that ps - c ESE(HAED S quantifies the contribution of the S to the feature produced by the FFN layer. Thus, it is the relative growth rate of egs that matters. Towards this end, it follows from (D.5) that, for all S E (H)S))(
0 log
4a (
c
-7xx(5t)
-
2 A
4a
(D.66
CD
Here we recall from (3.5) that ATx² quantifies the minimal gap between the modified x2-MI of S* and any other set in gD. The lower bound given by (D.6) ensures that for all S * S*, the ratio / grows exponentially fast, which further implies that ps* approaches one exponentially fast. This concludes the first stage of the training dynamics.
D.2.2 Training the First Attention Layer: Convergence of (w) One-Hot Vector
As we proceed to the second stage after ps+ & 1, it suffices to show how a(w()) converges to a one-hot vector EM+1-h for h S* in order to show that the model converges to the GIH mechanism. Recall that denote X = (#1,..., 2L) E Rexd For notational convenience, denote (hh) or (w) and let X(--M)(I-1) E RMxd denote the submatrix of X with rows
1 - (h) and similarly for _
(h) for any I, Following our convention, we let __ denote the (M + - --)tt entry of
Calculation of the Dynamics of (h), The main idea for analyzing (w(h) Jik= is the same as that in the previous stage: suffices to analyze the difference between the growth rates of different coordinates of (h) for E S* In particular, we care about how quickly w grows compared to other coordinates w is initialized to be larger than the remaining coordinates:
(h) 0,whh 0,
(A)
E
de
Ost
Osi
(D.7)
Osi
Ow
(A)
1= M
L
[xL+1 = = Ck)
Ös:
Ös:
E (as)
a
v(k)
Swnh
(h)
Ow (h)
M
Now, we invoke the result obtained in the previous stage that ps-
1, which gives us is R
Ihes (0) (h)
11t1) ). Consequently, for any / E S*, we have
Ost
II 40)
(k)
(K))
II
(0t (K)) o/21)) (eM+1-i-(o(0)) (h)
R3
ww:h
(K)
Ow
(K)
WES-
WES+)(()
(D.8)
where the equality follows from the fact that wh.)) only affects (u, (h) (111 and differentiating through the softmax function. Here we define bn = xa= + X(L+1-M-
(h) to simplify the
notation. Combining (D.7) and (D.8), obtain
0. (h) -h t z agT (o (CM+L-N
eM+1-1)++++ (OXAL-A--NMI1))) (1)
(h) Σ
_ (eM+1-A-- -
jfh
(D.9)
where we introduce the following notation
1 = ek) -1)
9h
Σ E oí(a s).
II
(0)
(N) (8))
y(k)
M++1
MES)(N)
30

---

A detailed deviation of (D.9) can be found in (E. Notice that o (h) (h)
(h) -o is positive at initialization. Now suppose --h - > 0 holds at current time L. Then, lower bounding 0,w (N) 0, (k) boils down to lower bounding 9h (EM+I-A EM+1-I) for + h. Furthermore, if we can show that 0,w"h d, is lower bounded by some positive value, the gap o-h o
(h) will further increase.
Since E a= = 1, this will create a reinforcing loop that makes o -h monotonically increase.
Relate the Dynamics to the Modified XMI by Approximations. We demonstrate next that gh (eM+1-A eM+1-1) for i * h admits a lower bound depending on the information gap AIy². Specifically, using the same strategy for (D.3), we have by definition that
T Ih.
(D.10)
1
Σ
E
1 (
x1 -1= EN) (k)
1) 1 = I141-5i0 e 8* {h}) by eM+1-i
L M
L=M+1
where for by we have by the same approximation of xi-h and 1/41 2 -1+1-1 as in (D.2) that
(h) by TTe41-;-- L+1 -1 + 1s (b)
-1+1-1 ( [x x1-)). (D.11)
Now we consider the case i = h and i h separately:
(i h) For 9 CM+1-A, we simply set = h in (D.11), and the indicator 1l (x1++11 = zl will exactly compensate for the exclusion of in the indicator function of (D.10). Drawing analogy to how we go from (D.3) to (D.5), obtain
and
gh CM+1-h R (S*).
(ii) (i + h) For 9h eM++1- with i + (D.11), apply the same reasoning as in the previous case. Additionally, by using the Bauchy-Sehward inequality, the following inequality holds up to small error (see Lemma F.7 for a detailed derivation):
In MEM4I-4 (8*) +
(8 ATx2.
Plugging this back into the dynamics in (D.9), conclude that for all # h.
0e (h) Dew (h) 2 a ( AÏx
(h)
Convergence of (w)())) Combining the arguments in the previous two steps, we can now say that (h) will monotonically increase. It remains to show that a converges to one. Note that log(o (h) /o (h)) (h (h) by the definition of the softmax function. Therefore.
0, log / _ -.
(h)
2 a. (h)
(h)
AĪ a A o (0) (00) /o (b))
where o (0) is the initial value of o TTh) at time t = O. One can now rearrange the term and pick the ratio 0 _ as the variable to track in the dynamics. A refined analysis in the convergence analysis in $E.3 shows that (4) converges to one-hot vector with a going to one. In particular, the convergence rate is determined by the information gap Tx according to the above formula.
D.3 Analysis for the Training of the Second Attention Layer
In the last stage, turn to the training of a given that all (h)'s for h S* approximately one-hot vectors. The following approximation the dynamics of a(t) is performed in the region v Ollog L), where the signal term in the dynamics dominates the approximation error.
31

---

Calculation of the Dynamics of a. After Stages and II, the output is approximated as y(k) 2 y" (k) = Et= of [ ek) for each k E [d]. Here the weighting coefficients oi,...,00 satisfy
of x exp (a s( X
Note that for each E [L], of indicates the importance assigned to the I-th token based on the corresponding history of zi over the information set S In the population counterpart, when the chain has sufficiently mixed, for given XL+1-S', we can roughly view each (æ), ((--) as being sampled from a reweighed version of the stationary distribution:
2r Xi-s (XL+1-S-) p¹ (x) X1-S-) exp (a ( 1(X--s- XL-1-S-))
Following the same argument those in the previous stages, replacing the sum over I with the expectation over the stationary distribution, we arrive at
1 - (= Cy X X
-
Oqa R. (E.N 2.2 2.2--55 yupr 1(x-s
7-8-) =7-8))
(D.12)
See detailed derivations the above approximation in $E.4. Comparing the above expression with (D.4) in Stage I, one see that here (z, x-ss) and (2, Z-s-) no longer independent because now the model has learned to perform a information-theovveic feature selection, i.e., focusing on tokens sharing the same set of features based on the information set S", which is defined according to the modified x" 2-mutal information. In fact, the underlying joint distribution q" is given by q²' = (x, X )+) (z, Z- X X-s)).
Divergence of a. As the dynamics of a has no closed-form expression due to the nonlinearity in the reweighed distribution fi"', we resort to providing characterization for cases where a is either sufficiently small or large. In both cases, the lower and upper bounds of (D.12) can be derived, respectively. Using these bounds, we can argue rigorously that for small a, it undergoes super- exponential growth until it reaches critical "elbow" value. After that, when becomes even larger, grows logarithmically until it reaches 2(log L).
E
Analysis of the Training Dyanamics
Masking the Simplified Model. Recall that we apply mask to the first M position in the simplified model. Therefore, we only allow index / run from M L in the following analysis. In the following, we first specify the conditions on L that are required for the analysis of the training dynamics and then present the proof of Theorem 3.6.
E.1 Conditions on the Sequence Length
We first introduce the following condition on L:
L Ω AT3 (1 - A 12(1-A)-447, VL Mvd,
(E.1)
where Q only hides a universal constant that does not depend on the model parameters. The conditions in (E. will facilitate our analysis for Stage and Stage II. For the last stage, we require
( (1- (112)1/2 +2trn/4 )). )
logn=! L 2 2M In --1
L (log L)a (HAnytaasTT
2 0 1
(
VM+* +d
(E.2)
where
::
(.1 (-1x-s.)) 1 (-)) A1,
and Ω only hides universal constants that do not depend the model parameters. Here, n² n= X X-s)) denotes the stationary distribution of the Markov chain over token x and its parents X-S+, with being the information set defined in (3.2).
32

---

E.2 Analysis for Stage
In this section, we analyze the dynamics of the parameters (C))SE(HTE in the first stage of training. We will show that there is unique S. E <D such that cg. dominates all the other eg at the end of the first stage. In addition, we will characterize how fast this happens and provide a corresponding convergence rate.
Proof Strategy. At high level, the strategy is to analyze 0 log cg. Ōe log c's for all S S S* via the following steps:
Dynamics Calculation. First, we calculate the dynamics of log ci for each fixed S. By selecting sufficiently small values for and e, and leveraging the mixing properties of the Markov chain with large L, the dynamics of log c approximately governed by the modified mutual information T (S).
2. Lower Bound for The Growth Rate. Consequently, we are able lower bound the differ- ence between the growth rates, 0, log es. - Ôg log c² in terms of AT,a, the gap between the modified mutual information of S* and the second-best set.
3. Convergence. Finally, we derive the convergence using the above lower bound.
Before presenting the proof, first remind the readers of few definitions and notations. Recall that our simplified model is given by
Esemign Os Ihnes (h) (M) L+17
y = (a(as)x)
Σ (asa) 2t, I=M+1
where
EseMM c
Also recall that Cp(()= EsE(HTE ct(() and ps(() = (s((t//p(tt) for each S E [H The loss function can be rewritten as
L E(q, where f= (- log(y +1))-
Here the expectation is taken over both the sequence (#),..., TL+1) and the Markov kernel P. We abbreviate a = (as) for convenience and denote by oR the (-th element of a.. By direct calculation, we have
OE Oy
= FL-1
Oy do
XT,
do = a o1((s) (eT - o),
y + +1
=
as:
Then applying the chain rule, we have
at de dy do
IL+1
-a
(xI- (a(as).
(E.3)
0s: Dy da Ost
y +
In addition,
Os:
2cs Ihees
(u) (h)
2css1
2cs
(oj (h) II KES
.0 $1)
Ōcs Ese(msD C
E SS &
CD
Now, ready present the proof of Theorem 3.6 for the first stage of training. We remind readers that here only {cs) {cs)s|H|g are trained, and we omit the dependence on for convenience.
Proof Theorem 3.6: Stage As discussed in the proof strategy above, we first derive the dynamics of log cg for each fixed S E [H] Then we compare the growth rate of cg. with any other eg.
Calculation of The Dynamics of log es. We fix a S E [A]&D and apply the chain rule 0E/Ocs = Ete 0e/0s, 0s//Ocs and the gradient flow formula that 0, = -2cs aL/Ocs. We have
0,03
4ac3
Σ 1=M+1
L
E o(asss)
( [IL+1] )
(x) (a1-w) 4)
II hES
(u) w - 81 )
-11)
(E.4)
Cn
y-
33

---

In the following, we consider fixed π for error analysis and take expectation over again when plugging everything back into the dynamics. To simplify the expression of 0, we define quantities 9o.s and f as
9o,s
Σ M=++1
Exia (1s)
1 = =1 =====e)
I(K)1(IL+1 = ex)
II kES
(ot (h) (h)
y(k)) +
u(k) +E
[1 = x1 =x)
u(k) ( = ex)
f f==
Ex|= 0i(as)
-8
(k)+5
kk) +
1=M+1
Note that here ff does not depend on S. Based on the above definitions, we can rewrite (E.4) as
1
de log c = - = Cp EEwplgo, - f
(E.5)
cg
Using this, it can be shown that (t) does not change during the training, described in the following lemma.
Lemma E.1. The quantity Cp(t) = Ese Cg(() is preserved along the gradient flow over (estse|Mss i.e. (,Cp(t)= O.
This lemma will be useful in the following analysis, and we defer its proof to $E.2.1. Next, we proceed further simplify the dynamics in (E.5) by approximating go.s.
Simplification of 0: log cs..
To approximiate go.s, we introduce the following quantities:
(xL+1= ex)
7(k) (xL+11 = =ex)
I1.S = L
Ex|s
II hES
M
[(kh +E
T(k) +
M=
92,8
1
M
Ex|=
(xt+11 =I1 ====NN ex)) -1)-
(=(())
II heS
(0) (b)
.02t1)
(=M+1
1(x = ek)
93,5 (2,,
10 (Ex) -11
II (M(0 (2), z (x)3
hES
where 7- = -(- is independent of X (x- (J:-...... x and we define
(1) Σ
(h) -- (Z)
(h)
i Z-in and y M
1
xt-
I= M++1
Here o(k) is the k-th entry of y. We remark that each of 91,s, 92.S 93.s is function of and but we omit the dependence for brevity.
From 90.s to J1.s, we replace attention probability (r(a)) by the uniform average with factor 1/L. which yields y. From 91.s 92,s, replace the empirical distribution with the stationary distribution μ² and drop the small constant E. Finally, from 92.5 to 93,s. we replace the average over the sequence by the expectation over the stationary distribution n² of the underlying Markov chain. We will show that the approximation error in each step is small, given that and are sufficiently small and the Markov chain mixes well for a large L.
For the approximation of go.s by 91,s, note that when is small, the attention probability ai(as) = 1/(I. M) for all l [L.). More specifically, it follows from Lemma F.3 that
Sad
Igo.s - -g..sl S e
For the approximation of 91 by 92.s. we leverage the approximation ((k) = n² (ex) due to the mixing of the Markov chain for large L. The result Lemma F.4 implies that
1g,s 92,5l s 4. (1 - 112 x² (mO 1l 11 1) + 2/M M -
L 11/2
34

---

where (-)-) is the initial distribution over the first In tokens. Here we abuse the notation of f= in yx (No (=5) denote the stationary distribution over the last In tokens. Since Mmin 2 T by Assumption 3.5, we have
Dxe(Holl=")) - -
(x)))
S
Σ 1/100 (x) (E.6)
Therefore, can further simplify the above bound
1 -92.s = o VL(1 A)yra (VII-->m :
Finally, the approximation of 92.s by 93.s follows from the mixing property of the Markov chain. In particular, it follows from Lemma F.5 that
8M + 192.s - 93,sl -
16 Dx² (1o 1=)++ +
Ln L(1 - X) 11/2211
L(1 - (11-2))111/2770/211
Combining the above results, and by the assumption that = a(0) =0(()))((//3})) and & = 1/VL, we obtain the following approximation error:
od
1
6 +
1
1g0.s 93.s| =
+
+0
LL1 - 1)
7
L(1 A)7 81+2+ra/2
1
s 00
VL(1 1)yr+t2
L(1-
L(1 - D/
VL(1 1)7
where we note that |S| S D for any S E [D and the last inequality holds by also noting our
condition on L in (E. 1) that L 2 2((1 1y- As a result, the dynamics of eg in (E.5) can be
approximated as follows:
0 log ds-
4a
EEwplog, S E, where |8| S O CDVLL -COVVTI-NTT t2
(E.7)
CD
where O()) hides universal constants that do not depend on the model parameters. Here and in the sequel, we let & denote an error term that is of the order ((a/ LC -)+=++2) where the specific constant hidden in O() may change from line to line, but does not depend the model parameters. In fact, we can show Cp remains constant by Lemma E. and a is not updated during this stage. Thus, the error term | is of scale O(aI-1//2).
Lower Bound for The Difference θ log Cg. θ log cz. The reason for approximating go.s by I3.8 in the previous step is that the latter more interpretable, in the sense that we can relate the modified x² mutual information (S). Recall that for each S E [H] the modified >mmutall information is
[<])) 1rt (== e (= (=
-1)
Ta( E. P. (2.2))mm
11 (2-
EEX
Note that / in (F.7) is independent of S, and will be canceled when computing à, log c... - O2 cg:
4a
2 log oz.. 0 log cs CD
EE-Plog,s - g3,s] 2/
Thus, it suffices to consider Ex w 03.s]. It follows from Lemma F.6 that for each S E [A]<D.
Ex-p [gs,s] satisfies
E (gs,s
II (o 02) (S)
[I-II(00) Ts²())).
heS
hES
This yields lower bound for E [g3,s-] and an upper bound for E,, - lla.s] for each # S* ie..
II (a ( I2((S)) -(1--11111
II (a (o.h))) (
E Ex~P(ga,s-]
Ix())))
hES
hEst
(1) (o (s) + (1- II (1))))))) (o (x)(S*), for all S + S*
S II
hES
hes
35

---

Consequently,
4a
0 log cs-
Oe log rg
Be 103.5' 93355 + 2|E
CD
4a
0 (2 Ix²(5*) - II (o 22 [x)(5))
CD
hES
hES
4a
(o (0_)
(o_h)) [x2(8*) - -2|E1
2
CD
NES
AES
4a
2
(o 12) -2 2 I (5²
II
(a (2 Aلx² -/
CD
hES*
hES
where the second inequality follows from the definition A7x² minss| -
(x)(
Moreover, since each (o (² 3 (0, 1), we have I 22 2== 2 ( Appling this to the above inequality, we obtain
(o_ for any S E !D]
O₂ log Es. Or log d3
4a CD
II I2(5)) + II (10h))² A 2Tx(57)
-
h=1
(E.8)
Exponential Growth of cz... We proceed to show that the first term in (E.8) dominates the error term & and thus leads to the exponential growth of cs..
Note that by Assumption 3.3, whh 2 w + Aw for all j # h and h E [M], where the quantity Aw satisfies
Aw 2 log (M - 1) - log ((++ AÎ
28 -1).
(E.9)
Recall that we not updating the RPE parameters during this stage, (b)) is fixed for all [H] So the gap condition (E.9) holds throughout Stage I. This conditions ensures that (h) >>>>> = so
In (a (h)) is sufficiently large. More precisely, given that head h is more focused on the (-h)-th
position by having a gap Aw in the initialization, we can further show by definition of the softmax function that
H II h=1
(h)
- (Mx1)exp(-Au)) 4h E |
(
(1 (M c)pp Aw) 21
(E.
Plugging (E.9) into (E.10), we have by additionally noting that T,=(St) 2 AT > 0 that
H II 2
AT
-1
28* + 2/3 AT
2
1+
141,²
21x²(8+)+ +
h=1
which implies that
H ( ix(")) + II h=1
AÏ - (8*) 2 ( AT 02 (o
(E.11)
h=1
Moreover, when L is sufficiently large such that L > l (( () -1-)) ), E in (E.8) satisfy [3] < BBadT2//CC. where Ω hides universal constant that does not depend on the model parameters. Therefore, combining (E.8) and (E.11), we conclude that
SaLix
a log e- O log cs Z
2| 2
(E.
3CD
2C'
This implics that & grows exponentially fast and becomes dominant.
36

---

Convergence of ps In this part, we treat all the model parameters as function of time t. For simplicity, we omit the dependence on when it clear from the context. It remains derive the convergence of ps- = fg-/CDn Expanding CD = Lselifee cs. we can directly calculate the derivative of Ps- as follows:
Cn
a -0(1
0, log(1 ps-) = O2 log
1
E SElHH|ED dg
CD 6
E
ESE(H]ED
cs.
(SSE|A<D (5) O + + EsE(H Oieg
CD
CD cs.
(SeEHH]go 2/)²
cs 0.es: cs. 0,d3)
CD(CD - -cg.)
SE(Blgo
Σ
2 cs (- log cs + θ log es))
CD(CD (s))) SE(M)
where in the last equality we use the fact that 0 log S * St, we further have
(0 Applying (F.12) to each
& .s. (- ad1x²
3 log(1 - -ps.) s CplCo c..))
Σ
2CD
sE(M|<D (88)
(CD- (-
adtx² C .a x 2CD
0.
CDlCD S.)
203
This implies that ps. = E /CD monotonically increases, and thus o3. (t) 2 ES. (0) for any / > 0 because Cp is constant by Lemma E. and eg. (0) is the initial value for eg.. at time = 0. Therefore, we can further replace cg. by its initial value the above inequality, which yields
2s (0)aSIx>
Ps (0)a
0 log(1 ps.) <
20
2CD
We remark that the above upper bound is independent of t. Finally, applying the Grönwall's inequality log(1 ps-). obtain
Ps- (O)aa :t
1 1- 1-ps- (t)
(1 - ps- exp
2CD
With training time t (2C'(0)) log () ps- (0)AI,a), we can guarantee that
1 - -sps- (41) < T.-1
This concludes the proof for the first stage of the training.
E.2.1 Additional Proofs for the Stage I
We conclude this subsection with the proof of Lemma E.I.
Proof of Lemma E.I. By (E.5), have
0 Exap ps(y0,s - f)]]
Moreover, by the definition of go.s and f. it holds that Ese(M PS90.5 = f. Then,
a,CD = Σ
d,d = 4a ETNP
Σ
PSSo.5 f = 0.
SE
Thus, the quantity Cp is preserved under the dynamics.
37

---

E.3 Analysis for Stage II
In this section, we provide the analysis of the dynamics of o (h) = (w for head E S*. For head h & S*, the results from Stage imply that ps + 0 for any S * S* Consequently, any head h & S* will be ignored when producing the output features FFN. Conversely, for h E S*, we establish the dominance of w w over w for all o h, yielding d_ + 1 as t - co. In this limiting case, head h exactly copies the (-h)-th parent. We also provide the corresponding convergence rate.
Proof Strategy. Similar to the proof for Stage I, our analysis for Stage II characterizes the dynamics of the difference between the positional embedding weights, 0,w 0,w (h) for all o via the following steps:
Dynamies Calculation. We initiate the analysis by deriving the dynamics of th for any fixed i and h.
2. Dynamics Approximation Then we approximate the dynamics by identifying the dominant term controlled by the modified x² mutual information I,= (S*).
3. Lower Bound for The Growth Rate By comparing the corresponding modified x² mutual (h) for all h.
information, establish lower bound on 0,w"h, 0,w",
(h) 4. Convergence. Finally, derive the convergence rate of d=h using the above lower bound.
Again, before proceeding with the detailed proof, we review the notations related to the dynamics of the positional embedding weights (w!h) } For the h-th head of the first attention layer, the positional embedding vector (h) induces the attention probability over a window of size M, i.e.,
a (w.h)) =: (h) (o _ (h) E IR 1xM
Further recall the attention scores for the second attention layer, as, where u U Then for each [L], the I-th coordinate of s given by
M
81 =
Σ SE(H]ED
II (uz
where each o7
Σ
(())) ((-MM)(I-1)
hES
Here ps is defined as in the analysis of Stage 1, and X(-M))(111> Mxd is the submatrix of X with rows M....,
By direct calculation, we have
(b(h)
Ou! (k)
Ow(( ("))(((())) (T)lh))T (())) E MRMM
XT X(--M:(--1 RRxM
do(h)
Then by chain rule,
(h) dul (b) (h( Ow)) (h) Ow(h)
XI_ MU-1 (ding(o - (h) To
RRXM
Moreover, we can view each i as a function of (o(t) ,..,, U he(M. Differentiating sy with respect to w/ and (h)
1 we have
Ost
Σ
II WES((()
(0) (1) (e E Ed.
0vi (M)
go sthes
0s
Σ
II (of (N)
e Rdx1
du
(k)
=
ps
L+1
SEiW)sD sthes
(ES)(()
In the summation, we only add those S's in [H|s containing h. Also, recall from (F.3) that
0e
FL+1
(xr v) (as).
(E.
0s:
el
Now are ready to proceed with the analysis for Stage II.
38

---

Proof of Theorem 3.6: Stage II. We start by calculating the explicit expression of the dynamics of (h) dew" ii and then derive approximation of the dynamics, which allows us further show the conver- gence of (f))
Calculation of The Dynamics of O,w(h) First fix h [M]. To simplify the notation, for each E [L] we define
by == (==((-M)(-- U + X(L+I-M): oi (b) R-
(E.14)
Note that w(th) is the parameters of the h-thh head and only enters each vi that 81 E ps Innes
(h) 1, ...., L + 1. Recall
(h) .'241>, and the RPE weight (hh) for attention head h only
(h) influences its outputs of and 1h+1 the sum. It thus follows from the chain rule that for each i [M], we have
0st
ast
dst
Ow
(1)
au
(h)
aw (K)
On (h)
L+1
II (E)
(u) (1))
(h()
1 x(-M)-(-1) (diag(o!!))) - (aah)) TT())) eM+I-i
(h)T
Ps
SE(M]D sthes
Σ
II
(v) w1tt) of ( (M)
X X(++1-M)- (diag(o - -(oll)) To eM+1-i
ps (ES((N)
S hES
Σ
II (E))(b)
(0) (1)
bJ ( -
(h)
ps
SD
SthES
where we remind readers that e, E R Mx1 is the i-th standard basis vector.
(h) Furthermore, along the gradient flow aww (N) = - 0L -L./Ow (0 -i it follows from (E.13) that
0,w!! Ex,x
Σ
Os Iw (b)
de Osi
a
E oi (as)
y + e1
IL+1
(xi (er-
Sw (h)
Os:
M=M+1
I=
11 [+L+1 = 21 Ex)
u(k) 11 (x = ex)
Os:
a
Σ E
(as)
u(k) +E
ukk) +
Ow (h)
L=M+1
k=1
(eAr+1-: (o(h)) (N)]
a EEnP
9h.0
where we plug in the expression of @s:/Ow!! (!w above in the last equality. Here the vector Ib.0 is defined as
gh.o
Σ Σ
EXI-
PSO
1 -
y(k) 1 ek)
II
(0) (K)
v(k) +5
v(k)+ +
MEST(N)
HAES
where 0i is the softmax probability for the /-th token in the second attention layer. Comparing 0,w
and @0 (h) we have
0, 0, a. BEAP
9h.0 (eM+1-A (CM+1-A-(0(0)))) (c)())))) 0 9h.0 T
T
(ea+tii--(0 (T(h))
(h)
(E.15)
Using the fact that End === 01h) = we can rewrite
(eM+1IA- - (o((1)) (1))
- ( (h)
(h) ("(((EMA-AA--EMM. + (0 (0 0 (h)
((M) (o(()))
M
=of(C04114- + (b) _AA)) Σ
o -j (cart1-h- -MM41-5), (E.16)
(1)
j=l
39

---

where in the first identity, we add and then subtract term o-i MM+1-N- Combining (E.15) and (E. yields for each i E [M/] that
(h)
ww (K) 0, (h)
(E.17)
E E 9h.0 00
(1) CCM+I----M411- (o -o
(1) E0 oy -CN-I-n--M-HH j=1
Simplification of 0. We proceed by deriving approximations to the vector 9h.0, which will help us identify the dominant term in the dynamics 0,w0h (h)
Ô (h) O, Specifically, we define
9h.1 Σ
Exia
oi(as)
(XL+1 21 ek) y(k)+ +E
ukk) (IL+1 = ylk) +E
ek)
II WEST
(o 2 2222001
M+1
1
Σ ExIt
[x1+11 =II
v(k) [21+1 = ex T(k) +
(U) (h) 10221000
9b.2 =
M
5(k)+ +
1=M+1
MES-((A)
9h.3 =
1
ExIT
1 = 11 ek) -1)
II NEST
(0) (N)) -
L M
(ex()))
t=
k=1
1(x = ======
II (u (2), o( (x))(( 2) WEST)(1)
9h.4 = (L,, ((e.
(ex)))
where Z = [3_
2-1 T E Mxd is an independent copy of X = =x B.--M...,,x111 E ERaid
and
(x) = Σ i=1
(h)
(2)
Σ
(h)
o 2-
2-11
i=1
(X,2) Z(u (k) (X)) ((v (Z)),
Σ Z
1 -
M
1=
The strategy of gradually approximating 9h.0 by 1, 2, Th.3 and 9h.4 is similar to the analysis in Stage I. To see the intuition, from 9r.0 to 1 S, which is a result of Stage 1. From gh.1 to Jh.2, we replace by the empirical mean y. thanks to the fact that (a) x 1/L when is small. Then, from ga.2 to yh.3, we replace the empirical distribution y with the stationary distribution of the Markov chain. These two steps also appear in the analysis of Stage 1. Finally, to go from 9h.3 to Jh,1, we leverage the rapid mixing of the Markov chain.
the fact that ps+ = 1 and ps 0 for any other
Note that the common structures in (E. are IT - -EM+1-5) for i * h. Hence, we only need to understand the approximation error in each step for ST, (eM ear+1-)). Recall that we are focusing on h E S* in this stage.
From Oh.o to I,.1, we remove the terms the summation that are weighted down by ps for any S j S* due to the rapid dominance of ps- from Stage I. Recall that ys- converges to one an exponential rate while all other ps converge zero. For simplicity, let define
1 [ = x1 - Ex).
y(k) 1(1141 = Ex)
P(S)= = Σ Ex m (ass))
Σ
u(k) +
v(k) +E
M=
II
(N) (K')
, "00.2211000 (eNs1-n -M++-i))
MES
40

---

By the triangular inequality, we have
[gnoo 9hh11 (eM+1-M- -cM+1-1)
Σ
ps pSS)(S) p(())
S ps.) ( |p(5*)
Σ
Ps (o-1((S) 16(1 Ps*);
cE(H)
where in the last line we use the claim that [p(S)] 8 for all S. To see this point, note that by definition of by in (E. 14), we have
bb (eM+1-- -M41-1) = (v
WLL1 -i--) - (v) FILIIA - xr+1-1)| S 4,
II WESMIA
(0 ( 1
(h))
since (0 and have norm most 1. Then, by Lemma F.2 where we plug in the upper bound for the function f(-) in the lemma, we conclude that p(S) I 8, VS E [H] <D \ {S*).
Define A := ps- (f1), and A 1/L by the results from Stage I. Thus, obtain
I(gno - 1h.11) (EM+1-N eM+1- = 1621 & 16/L.
For the approximation of gh.1 by 9h.2, we use the fact that 0i (as) 2 1/L when a is sufficiently small. Specifically, we also take the absolute bound for f(-) as 4 in Lemma F.3 and obtain
32ad
(9h,1- g -MM-I-0)
-
g2
For the approximation of 96.2 by 9h.3. we the fact that 7(k) n²" (ex) ) for large L. More precisely, it follows from Lemma F.4 with the upper bound for f(()) in the lemma that
((9.22-hh.) - 9h,3) (EM+1-A - -MM+1 & 16 (1 (1= (po |l p²") + +
47 -1e.
11/2
Finally, go from Gh.3 to Ih, leverage the rapid mixing of the Markov chain. Intuitively, when and L ++1 1 are far apart. and its parents in S* are independent of æ and its parents St. This observation yields the approximation 9h.3 (CM-1-h eM++1-) by 9h.4 (EM-1-1) eM+1-)). To simplify the notation, define two scalars
d 11 (0r+ I1 =x==
:h.1==
II (ES
(o
2410
(8)
1 (ex))
0h.4 1(x ck)
II NESt((K)
(u(h') ( (x))))))
n (ck)
Using the notation above, have
[inh,3 - 9h.4)
(On,A)((EEE)))))) (EM+1-A--MM
L
Ex |
(1. our [ūn, Carti-n--MaA
I. M
1=M+1
Recall that
bf (eM+1-A- - -M---:)
(1 11tt 21-5) - (u. FL+1-A - -I1+1-)),
(b)
6(X, T
(b()) こー - (Z),
(Z), x-h -x-i)-
41

---

We apply the triangular inequality to obtain that
[9h.3 1 gh,a)) T (CM+I-h- - -Ms+1-1)
1
$
1
Σ
Exts [Tn (on
Bee. ( (Z)
(Z); x -iil
L
M
1=++1
Σ
Ex
(1)
Elx, Th (Z), x xلام
M
1=M+1
+
1
(1) Σ Ex|s (0t
(X); =n)] , ,IL+1-N1 (Io..)) ( (x),
M
M=++1
Σ
x1+1-11] Ex ( - -B(,,)),( ( 21
+
M
M=++1
Each term on the right-hand side can be bounded by Lemma F.5, where in the lemma take (o(1)) Jwes. E M and ((alk') ((o"))))ess ()ss.()
RMXISSI the two lists of vectors the M-dimensional propbabbli simplex for o and respectively. Consequently, we have
(A))))))))))))))))))) e4)
8M
16
(gh.3 a (EM+1-A- -MM+1-5) S
+
In
L(1 1)/1
Combining the above results and setting = 1/VI, a= a(0) s 0(1/ and together with the conditions in (E.1), we have
[gn.o " CCN-I-A-eM+1-11 0
(V(VLL L(1 - +2
where O(-) hides universal constants independent of the parameters the model. We remark that while the left hand side is function of t, the upper bound is independent of Then, we can rewrite (E.17)
0, "! (h)
(h)
M
= a ETAP
0 9h. - (1041-A-MM4110 (o - _
(h)
Σ #=1
-j (eeatt-n----HH
(h) T
( (b) (a
(h)
(
Σ
a 2 (8).
(E.
j=1
Lower Bound for The Difference aw (h)
0, ". To show a,w
(K)
(b) 2.w
(h) > O. we first derive
the lower bound of Ea~plghh 19h.4 (EM+1-h em41-=)]] for any o h. Since (ar, X) and (z, Z) independent and identically distributed, by the definition of 6(X, 2).
En~P [uTTa (EM+1-A -M+1-=)]
1 1 2= ck) -1)
II (2))) (X))(X)) (x), (X), zni)
1 (Ex)
WES
1(x ====e0
=ex) ex) -1)
II (N (n)) (x))(()) (x), 2-1) =
-2
1 (ex)
MES
2Th.1 -2Th.21
42

---

where we introduce the following quantities for convenience:
1 = 2 = Ex)
-1)
II (ulk) b()) (x! 2-11
Th,1 == (e..
(ex)))
(Es
1 1 [ 1 (ex)
Th,2 = E (e., Onr
II (2), uZ), (X)) (x))) x
ME
The quantities 1 and 2h can be further approximated. Specifically, by applying Lemma F.6 to Th,1,, where in the lemma we take (o())) IR Mx|"" and (wrls)) (h)). en) Rxx1 as the two lists of vectors the M- M-dimensional probability simplex for a and Ø respectively, and we obtain
(K) II (2 o Tx((S) s (11-
II (2 (N) (x)
Th.1
(F.19)
(E)))()))
MESt((K)
Drawing on the analagous reasoning as in the proof of Lemma F6,, we can approximate Th,2 as follows:
Th22
II (E)
(2 - 1 s (1-
(h)
<11-
II (())
(o - ) I2((8)
(E.20)
where
1 = ex)
3 == (1..
II (ES
(x=== = =-w) 1(x- ==-))=
10t (ex)
To establish the lower bound for Th,1 - Th is approximately equal to 71,2- We invoke Lemma F.7 with S = S* and S' = [ U [i] in the lemma to obtain
, let us begin by establishing an upper bound for dh which
1 + 1 Ix U (i)) - [x(8*)
1 2/2
& S 2
vi Vi h
2
Leveraging this for (E.19) and (E.20),
2Th.1 2Th,2 Z
II
(a (2'2 - -hhh AIxx -a11- 11- (ES+\(A)
(h)
II
(o (o (h)
--hh Ix*(
(ES+((A)
((h)
II (0_)) AĪxa -4 =(11 II (n)))) (S*). LES
(E.21)
NES*
where in the second line we multiply an additional o o_A
to the product as o o E [0, 1].
Next, we provide lemma showing that (h) -h is growing for all time 2 t1. where t the starting time of the second stage.
Lemma E.2 (Reinforced Growth of o (i)) For all h E St, we have for all + at any 2 t1:
0
(h) O, and o log (h)
0 log 8hh)] 0, 0r (h) O.
(E.22)
Proof. See 1 for the proof.
In the proof of Lemma E.2, will use the following useful proposition.
Proposition E.3. Suppose Thes. ( 2 1/(1 + (M - 1) exp( with Aw satisfying (3.6), and -h o for any & h St at a given time t. Suppose Assumption 3.5 holds and
(h)
(h)
satisfies (E. 1). It holds that
0 log σ -h 0 logo (1) Dz (h)
(h)
-.,
(M)
(h) 0= +o
(N)
Σ j=1jjh o > 0,
0,0 0,001) (1) > 0,
di thE
(E.23)
43

---

Proof. See $E.3.1 for the proof.
Lemma E.2 implies that during Stage II, for all + h and E S*, we have (h) wht and o > - for all t 2 t1- In addition, (h) is growing, all the conditions in Proposition E.3 are satisfied for any 2 f1, and hence all the conclusions in (E.23).
Convergence of (A)). Finally, we characterize the convergence rate of g(h), For the convergence analysis, we adhere to the convention used in the previous stage, treating all model parameters as functions training time I, where = 4 marks the start of the second stage. With a slight abuse of notation, we denote by (h) (t) the value of a (=i(w(()tt)) at time t, where w( is the input to the softmax function, and a-)(-) refers to the (M + 1 - i)-th element of the softmax probability. For simplicity, sometimes omit the time index when the context makes it clear.
Note that 0:0 (h) -h for all h St. Hence by the definition of the softmax operation, we have (N)
σ (h)
0 -h exp( (pp(-(w-h w (1 2 a (t1 exp(- (b) w (k)
(b) - (0) exp(- (w w ).
(E.24)
(h) and the second line follows from where the first inequality follows from the monotone growth of d the fact that the first attention layer is untouched during the first stage. Note that here in (E.24), (h) and (h) are functions of I. Now, putting together (F.23) and (F.24), and also noting that Ohh for all i # h and E S*, it follows that
(h)
0e (k) - D, a (h)
o =A exp( exp(-(w_th w ()))
2
2
Rearranging the terms, and using the fact that w
(t1
(b)
(tz)
w (tz) 2 Aw by Assumption 3.3, we get
(0)
exp w (t) w
( 2
(t - t1) + exp(Aw).
2
This yields lower bound for o ( (t) follows:
(t) 1+ I exp(w (h) (t) w
(h) (t))
1+ (M- 1 min (0) (t t1)/2 + exp(Aw))-1
where we define min (0) = minies (h) (0). Consequently, we have
2/5"1
1 II (o (0))))) 21- 1 1) (aSie mii (0) (t t1)/2 exp(Aw)))
hES
(1- (1- mmin m0nn (t) exp(Aw)))))) + (M-
(M (M-1)
2/5"1
(aDix)
Now, we consider large / such that
(M - 1)
(adix² Gmin (0) -t t1)/2 + (w)) + (M - 1) 2|8*|
(M-1)
Then, apply the inequality (1 - z²" > nz for [0, 1/ni] and to obtain
2|S*| (M
1- II (h < aهL mmin (0) f2)/2 exp(Aw) + (M 1)
REST
Therefore, with training time to S (M - 1)/aAI,² mmin (0) f1, ensure that
I BES
(o (K) ( -.-
This completes the proof for Stage II.
44

---

E.3.1 Additional Proofs for Stage II
We conclude this subsection with the proof of Lemma E.2 and Proposition E.3.
Proof of Proposition E.3. The condition Ihees- (o (o_ 2 1/(1 + (M - 1) exep -1)exp(AAw))2s with Aw in (3.6) implies that
+
II BES.
(o ( 22 1+
(E.25)
141x=(5-)
(5²)(S²) + ATx²
Combining (E.21) and (E.25) yields
Enpp 9 (CM+1-h -M+1-i) = -Th,1 - 2Th, 2 A
(E.26)
3
for any + h. Applying (E.26) to (8.,8), since each o (h) > 0 and (N)
(b)
o-h
> o 0-s at time / for all
i + h h E S*, it holds that
0, (h) -,w
(b) 2 a a(o0h) (o -(0)))
Σ
o (ZATT= -1e1)
Then since we assume sufficiently large L N A((A) P2. (1 2) holds that |8| s
AĪya /6, we further have
aلI/² 0, (h) 0,w!h)
M
- + (o -
(h)
(k)
Σ j=
(لh) ا
>0, Vit Vh
2
As Og log (h) de log (h)
= 0, (b) - d,
by property of the softmax function, and > for all h E St. This completes the proof of Proposition E.3.
E 0,00 (h) 0. we have 0,0 (h)
hh
Li
Proof of Lemma E.2. We give proof to Lemma E.2 by contradiction. Note that at the beginming of the second stage / = I1 we have all the conditions for Proposition E.3 satisfied by the initialization
(k) conditions Assumption 3.3. Then, by (E.23) Proposition E.3, we have Ū log o Og log (h) 0 and 0,0% (h) h > 0 for all + and S* at t1. 11.
t is the smallest time such that at least 0,0 - I 0 or O log o ah) O logo"! (h) < 0 for some i + h and h E S*. By definition of r, we have (E.22) holds for any moment E T). As o ( and the gap a - o - are monotonically increasing, we have by the initialization condition and the boundedness the gradient that time T:
Next, assume that
II (0 1/(1 + (M - 1) exp ()))((A)))) and KES-
σ o
o o Vi h, th S*.
(h) Hence, by Proposition E.3, we have 0 log ( - O log a > and 0,0 > 0 for all i + h and E St at time т. which contradicts the definition of r. This completes the proof of Lemma E.2. O
E.4 Analysis for Stage IIII
In this section, we derive the dynamics of the second attention layer's weights a in Stage III. We characterize the dynamics of a when < O(log I.). where the signal term of the dynamics dominates the approximation error. We provide the growth rate of the weights for two regimes; when a is either sufficiently small or large.
45

---

Proof Strategy. We analyze the dynamics of a via the following steps:
1. Dynamies Calculation. First, derive the explicit expression for the dynamics of
2. Dynamics Approximation. We approximate the dynamics by exploiting the mixing proper- ties of the Markov chain and the convergence of the weights from Stage I and II.
3. Lower and Upper Bound for The Growth Rate. Finally, we establish the upper and lower bounds for the growth rate of when is either sufficiently small large.
For set S S [M] denote X (#i_a and simply use X-s. In this section, we abbreviate ps ps(4)) after the first stage's training as ps, and o (t2) after the second stage's training as o _.A
8). If O, will ignore I in the subscript
Proof of Theorem 3.6: Stage III. We start with the explicit expression of the dynamics a.
Calculation of The Dynamics of a. First by the chain rule,
de Da
Σ
de (ası)) Σ
I.
I++1 w+el)
(x1 (11-y)
y) ai(as) st.
M(ss))
au
1=M+1
I=M+1
where in the last equality we remind readers of the same procedure we have used in the derivation of (E.3) in Stage I. Then, taking expectation with respect to X and π and expanding
81 = EsElllg Ps Ihes (v)") .224)) we have
0,a=
OL
E
ØL+1
(z) v) o (as) 81
da
D +
M=
1 = 21 10==00 cx) y(k) 1 = cx) y(k) +E
=E
Σ
PS
II hES
(0) , L-11)
=: fo
2kk) +E
Im
M+1
Jen
We remind readers the shorthand (as). We denote the above quantity by fo-
Approximation 0,a. Similar the analysis for the previous two stages, we develop a sequence of approximation steps that transforms 0ge into tractable quantity. We aim to decouple 2L+1 and Iı, approximate 81 by population version, and transform the expectation to one under the stationary distribution of the Markov chain. Specifically, the approximation involves the following steps:
Our first step is to remove the summation over [/]< (SS* where S* is the optimal set that maximizes the modified mutual information defined in (3.1). This is because cg. dominates by the analysis of Stage I. Specifically, we define
f f==B
Σ
(#L+1 xq ===e
v(k) 1 (#L+1=
II
(of
(h) . L+1
bkk) +
u(b)+ +
M
To bound Ifoo fil, note that for any S E [M] since each di (h) has norm at most can invoke Lemma F.2 with C = and obtain
[ x1 ex) v(k) 1 = =ek)
II NES
(M) , 2411
(K) (8)
$
y(k) +E
y(k) +
I=M+1
k=1
It follows that
d
1(1: ex)
y(k)
Ifo - fil = E
Σ ps Σ
1 1xIL+1 = ex)
(uf
y(()) +
v(k) +
SS(H M8+ 1+++1
k=1
hes
L
(1(m)- = ex) 1(k)) +
y(k)
+ - (l-pss))
E
Σ
1(x1+1==00
II
u(k) +
4=M+1
1=1
hES
$ 4(1 ps (t1))
24,,
where ps- (t:)).
46

---

In summary, the difference between fo and fi is controlled by the convergence results from Stage I.
Our second step is to characterize the approximation error incurred by the difference between the ideal attention scores and the actual attention scores the second attention layer. Let us define 57 = - 1l (21-1 T1+1-A) the ideal attention score for the second attention layer. We invoke Lemma F.1 to have for all I E [L],
|s| si v A + A2,
where 22:=1- - II
(o (
LES*
Corresponding to (s we define
exp (a Ihes. 1(=1-1= -ILIIAh) of = MPT exp (a lall (1r-- x(+1-N))
(²(k)=
Σ of 1 I=M+1
e
In the vector form, we have y² = Ei=M41 oj Leveraging the above approximations, we define an approximation fi as
(( (*) ( Ex)
f2
E
Σ
II
1(x = -(1=-1=21++-)) x1+1-1)
y (k)
(²(k)) +e
I=M+1
1=1
AES+
Applying Lemma F.9, it holds that
fi f2l S 12 (1 a(t) 11 (A1 + A2)
In summary, this error terms captures the difference between the ideal weights and the actual weights obtained by gradient flow at the end of Stage II.
Note that y" (k) also random due to the randomness in oj, and as I. is sufficiently large, want to replace y (k) with its population counterpart. Let X and Z (z- (2-M,..., E d be two random variables and we define similarly for E X and X = (I-M..... x-11) 11 E x,, To this end, we define a reweighed distribution
1 (2, Z) exp (a Ihes 1(2- = =x-n)) 1et (2, (= Ex.2 pMT (2', Z') exp (a IInEs 1(<L_, = x-))))
(E.27)
where nM is the stationary distribution of the Markov chain can be viewed as a reweighting of the stationary distribution over (2, Z) by an exponential term that depends the sequence X-s.. We use pit (a ex X )) to replace y (k) and define f3 as
window of size M+1. This
1 = 21 =1= ck) (== eg XL+1-S-)
/3
Σ M+1
[ = ex)
II hES
1 1
1=
One immediately draw connection to Lemma F.4 both targets characterize the gap between the empirical and population distributions. The only difference is that this time we have the distribution reweighed by some exponential term. For completeness, we provide the approximation result in Lemma F.10. which bounds the difference between fz and fs as
ff2 f 8(1 8(1-A) 1/2 xa (No ll p²) ) + 1) + +8VM
2dE
VM + d
11/2 7 |5|+11
Y
11/211
where (-(() is the initial distribution for the first r. tokens in the Markov chain. Here and in the sequel, we simply use Dx² (Ho || (²) to denote Dx² (po( Ksr.. =.) || 1²" (Xirr, = when is clear from the context. In the last inequality, we use the fact that Dx² (1o p²) S - n by (E.6) and the condition = = L 1/2
Note that the expression of fa, each of still implicitly depends on the actual value of the sequence X. Since L is large and the Markov chain is well-mixed, approximate
47

---

Σl-m+1 of ((-s))))))) = (3-)) by frt (.-11141-85) This gives rise to the following
approximation of fs:
1 === [xT+11 = =ex)) -
f4 (x= sr)
1100 (2 = Ek X1++-5+))
11 (x ck X pr (z ck2 Z- X-s. (x-s
(= = En X x
pT = X X-s))
Applying Lemma F.11 yields
87 (1 (1-A) (12))) ((*) + 1)1/4 + 16y 16y-1 'VM
VM
If - fal S sup rEsupp(P)
L 1/181+1
L 15//2//111
where the fact that Dx² (mo ll 1₂") S 7h by (E.6).
Let (z, 2) ~ (-1 Since is large, the distribution of (æ1+1, X1+1-s)) is close to the stationary distribution μ². Thus, we introduce the following approximation of f4:
fo = (0)))..) (= (+)
1(x ======a ex) Tr (Ek x-s))
11xx = =ea)
II NES*
1(2x = x-N)
n (xr ek X 5t ) (z Ck.
Z- X-5 (x-s))
(x, )so)
1 = Ek X-s))
-B- = X- 1 (--s.)
(E.28)
Note that
(x ck 1x c, Z- ss x-ss))
(= Ek |
d
Σ n0 (r = ex X-s)) 1F" (2-s = X-s X-s 100;==e00 ex) s
Σ xº (r ex X-s))
=1,
k=
and so is (Z-s- X-s. (x
The difference between fa and fo is thus
bounded by 2|pp" (I X X = -)- - 18t (2 X == l) and by the
results in (F:29) of Lemma F.16:
X--M
- fols6
sup TEsupp(P)
L M (40 +1 g/ol S L-
where we use D (Ho ll (²)) Tn and the condition on L in (E.2).
Collecting all the above approximation steps, we obtain (where we use S to hide absolute constants)
VM
ffo fsl A +1+ (A1 A2) + +
11/2
VM + d
S a. 1 1/2 + L1/2 15-
48

---

where the last line holds by moting that with sufficiently large tj and t2 we have A + 2 S L-! and L 1-122 Here, express the error in terms of the trainable parameter and define
VM + d
{(a) x L1 2)1/222151+22 + L 1-12
In particular, we have for = (log L) that
VM + {(a) = 0 11/2(1- (12 L
log I.
(E.29)
(+2tr,/4
In nutshell, we conclude that when the weight a satisfies O(log L), the dynamics of can be approximated by
0.a + E(a)-
(E.30)
The following proposition helps us reformulate fs in a form that facilitates the analysis of dynamics of a.
Proposition E.4. The term fs can be reformulated
Js - [[(X J(X-s+;a,m) E. (x"(X-s))) 1 (X (X-s)).
where r" ("=ss=;;) = (1 + 1¹" (x (6" _ 1)) is the inverse of the normalization factor of jr" in (E.27) and
(n (r es l X-s+) - (x = ex))²
J(X s
(X (x=)+;a)) 1 (x ek x-s)) rrt (X-s+;a) (0 ex)
kEld]
Proof. See $ for the proof.
Inspired by this form, we define an alternative function J(:rr, x)) as
(1r( = (x= X-s) 4 2x))²
-(x-s-;7,))= (1-r) (x = ex X-s+) + ( (=
TE [0.
(E.31)
kEld
where we replace r*t (x= by parameter r E [0. 1]. As exactly calculating the inverse nor- malization factor nM (X-s+;a) is intractable, we instead seek find an upper and lower bound for (X-s+::) and plug them into J(-ir, TT to bound fs Suppose that r (X--ss-ii enjoys the following parameter-dependen upper and lower bounds:
(a) S r* (X_sx; S (a), VX_s- E x1s") VI suppp(P)
Thus, an upper and lower bound to (K-ss;a,m can be given by
inf
J(x ((;-s:; S (X S sup 7(--s:;7,n re(r
(elr-())r440
(a).r) (o)
In order to effectively tackle these bounds, we then study the properties of J(::T, x) next. Proposition E.5. Define
(x-s+ max (D (1)
The function J(x-s+; σ) with r [0, 1] defined in (E.31) satisfies the following properties:
1. J J(X-s+ir,,) is convex in
2. J(X-s S D+(X-st
3. J(x ((x-)+;,,,) is Lipschitz continuous in r with Lipschitz constant "D (X-s...))
Proof. See $E.4.1 for the proof.
49

---

Upper and Lower Bounding J(X-s+i a, T). Previously, we show via a reformulation of fs that it suffices to bound -(X-s+; a, σ) In the sequel, we let
P S., m) = max { {Dx (u ()) 1 X-s))),
Dx (N (-10 (-) (1)
D (X- 5t,
D D+(X-s+,7)
max max
x- Dxe(pT AM (-
-- Dx2 (1 (-1xss))) | (T))
max
can be noticed that
Dyplu
Dxa(u X Xx (1)
max
- Dx=(1E (-1xss)) (.)) x-sr,= (u ()) || 1r (. (.1X-s-))
max
107 X- 1.TT (-1(--xs)))
10
ST
max
max
max X-8+
(.)
where the second incquality follows from noting that the x2 divieeegee defined as D2a = Ex (u(x) (()))//((>)) and Dx²(" || xx) (v μ) S sup, u(x)/v(x).
Apparently, r= (X-s+: a) is function of and enjoys the following parameter.dependent upper and lower bounds:
r4 (a) (1+ min p= (X-s))(e" 1))=
x s
(a) = (1 + max nº (X ((e -1))- x-
If a is small, we see that both r+ (a) and r (a) are close to 1, and we directly have
(a) S r" (X- S 1, where 1 max x-
p² (X-s 1) ra) < 1.
This suggests an upper bound of J(X 7(-=+;;am) < sup
as
J(x-s+ir,m) J(x st 1, +7- D+ D+ R)) : (1 (a)r-(a))
(a)))
< Dz (-1 11 1x" (1) + D+ (T St
1e (X-ss)) (E" 1)
s (x (,)) 1 X_
p² (X-s.) (e -1);
where the second line follows from the Lipschitz continuity property, and the last line holds because the ratio D+ (X-s+, Dx² ( ) (-)) is upper bounded by p. and further by 7 A similar lower bound can be obtained by changing the sign y -2 maxx_ 1t s-)+)- (e 1). Hence, we h
V(X-s+;a,t) = (1 (1+77,
nº (x-s.)) .(e" (61-1))
x s
(F.32)
On the other hand, when becomes large, we have both (a) and r- (a) close to O, and we have
0 < rx S (a), where 0 < (a) <
minx s. m 11 (X- )(Ea -1)
In similar fashion, we have the following upper bound:
(X=s+;a,) S sup
J(X ((X-s+;r,)) < (x;,5= + D+ (S,- (a)
rE(0.r+(a)]
Dx
=2e(((() (xX-s) (
minx
nTT (X- )ee 1)
x-2
< (e
1 minx s p. (X-s+ (ea 1)
We can similarly obtain a lower bound by changing the sign of the second term inside the bracket. Hence, we have
=(x=s;;a,))= = (-))-))) (1 minx sa 1a (X" (ca - 1 (E.33)
--
50

---

Divergence of a. Recall that we have shown the dynamics of a in (E.30), where {(a) is negligible when goes infinity. Thus, when L is sufficiently large, we see by the nonnegativity of fa that a(l) continues to increase as increases until reaches a point where Ss no longer dominates the approximation error. To characterize the regime where fs E(a), we first note that for I log L it holds by (E.29) that
{(a) = O(L- log L) = L-1/2
where z hides logarithmic factors. For fa. we recall from Proposition E.4 that
J(X-s+).e
fs = E ox gomp (1 + (X (en (en-1)) 1)) 1 (X (X-s+)
where for small we have /ss (11) and for large a we have fs lie (.)). Thus, e > L-1/2 gives the condition for fs to dominate the approximation error, which gives a = O(log L). In the sequel, we consider the dynamics for a S (log L)/8 88 and give more rigorous analysis.
We the notation x o(1) denote that term is much smaller than 1, for example, (log log T.)-1 = o(1). For any zo and 6, we write = + 5 to indicate that æ is bounded within [wo 6, xo + 6j. In the following, we assume there exists 8 satisfying S 72/4 /4 A 1/8 and
8. HRN
Σ x S=
(x (11X (H (X- (X_s:)))
2 {log L),
Dx((*()) (-1-s)))))))) LL 1/47 1-1444
8. E ETpP
2 E(log L.)
11 (X-s.)
g+
Note that
Elog L) 5
VM + d 1,12(1 1)1/2 8-1-2-fa/1
log L + T,12
By additionally noting that n²' (X-s-)
2 115 thanks to the lower bound of the transition probability,
we are able to find such a if we have
L
1) 8112tr/l )")
VM + d
2
(
(log L)1
where k is defined
R k: (D=
E (1 11 (-)( A1,
and 2(-) only hides universal constants. Note that this is already guaranteed by the condition on L in (E.2). In particular, we can just take 5 = >² 2//44 A 1/8 in the following analysis.
Small Consider the case where is small in the sense that n² ("(Xx-) (E" 1) S 8 for any X-s- and π E supp(P). In fact, one can directly deduce from our previous results that 8 s (a) I rd (X)))aa) < and
1- 35 S (x=saaa))))) S 1.
For J(X sx .T), combine the condition that 1z8 (X_s.) (6 11) S 8 with (E.32) to obtain
that
J(X-s+;a,1) (1 (1 Dx² (- xI where 7-28 < 1/4.
Combining the above two results with Proposition E.4, we have
fo -B.X_ [J(X [](X-siia,,) e (r" 1h 10 (X-s))
=11 (1+(7- (7- + 3)0) EraP
Σ
Dze (.1X-s)) 1² (X-s))) .0
X_s.
51

---

Also, the noise term ç + (a(a) is upper bounded by
E (log L) S 6 Enup
Σ (x 1 (X_s))))
X-xe
S 6 Enap
Σ
1 (X-s:)2 -
by the construction of 8. Combining all the above results, we have the dynamics of
3,a = (1 (x~ 4)6) Exap
Σ
(1X=) μ² (X-s:) E.
X-
A simple reformulation gives - a (1 (1= Enup Σ (=)
(x
which implies that for small the growth follows
a(t)= & -log
e (--a(0) - (1 (x-2 + 4)6) Exup
Σ De((T((1xss)))1 (1)) sx-s :t) X-8*
Σ [2((((( .1).
a(l) 2 - log aa() - (1- + 4)6) ETOP
Therefore, in the beginning, att) grows super exponentially fast.
Large a. As a grows large such that μ¹² (X-s)) - 1) 2 8-1 for all X-s. and π E supp(P), conclude that 0 < rd (X= $ (a) S ô and
(1+ 11+w0(M (X ( - 1)}3 = (1- 10t (X-s-)
(= (X s. (x--ss)ee))
10tT 10 (X- ( 1) )²
which imples that
r ax (n -s (X )oa) -3 1.
1-
For (XX-s* a, T), we combine the condition that 1t" (X-s:) (6" -1) 8-1 with (E.33) to obtain that
J(X_s-;a,)) = 21+ Dzs(E" (()(((()))))) ()) ))) where 2- < 1/4.
Combining the above two results with Proposition E.4, we have
fs = Ex [](X s-;a, (r*(x=):)) 1 (x
22aa
= (1 (2- Eaap
Σ (1 (x X_
For the noise term & + (aa), we have
-2a 1 (X-ss)) 25+b(a).
8 E ExNP
which can be verified by the condition on ò as well as the fact that we are only considering a S (log L)/8. We thus have for the gradient that
2a
aaa (1+ + (7-2 4)0) EnNP
Σ X-s-
Dx* ((x-s))))
-=(X-s)))
52

---

By rearranging the terms, we further have
2ee20 (1 4)0) BRND
210 (X-s)
Suppose this large regime starts at to with value a(to). Thus, for large a, the growth rate is characterized by
a(t) = log (1 + 4)0) Exup Σ
(t- + eza(t0)) )
2
(x
x_
which is logarithmically fast. This step ends until reaches the value (log L)/8. This concludes the proof.
E.4.1 Additional Proofs for Stage III
We conclude the proof of Stage III by providing the proof of Proposition E.4 and Proposition E.5.
Proof of Proposition E.4. In this paragraph, we aim to gain more insight in fs. By the definition of fs in (E.28), we can rewrite fs as
1- = Ek
X x- -1) (2 = xx=s - (x-s))
Js = E X- 8 mup
T.Sh (z ck X ss)
n (x ck x X-s+) -1)
?
E gumpp
(F ck X-s)) ( = X 1x-s
(z c x-s+)
where in the last step, we use the simple fact
Σ M Σ -1= Σ
p(X x - -1) 2 q( x
qX== = xxx=xI>>
In the definition of Ss. the key quantity we aim to understand is the reweighted distribution Trt (z, 7 | X-s)). For the readers' convenience, we copy the definition of the reweighted distribution
here:
1r 2) exp aa Ihes 1(2
N (z, 2 X
(E.34)
Ex 1. (=', Zz) exp (a llnes- 1(2',, = =x_n))
key observation is that the reweighting only depends on the value of Z-s.. Let [M] and denote by Z-s- = (z~n)hes.. Following the above observation, we can additionally condition on Z- and conclude that
7rth (2, 7 5., 7 1 x-)
xrd (z, 2s (=s.,, = Ex 8. (=', 7..77.7. X-s.)
n² (z, 2_, 5:) exp (a Ihees 1(2m = x-)
Σx 8t nrt (z', Z' s-, Z=s.) exp (a lllesss 1(2 = x-n))
11th (z, 7, 7-5 1x7 (Z-s))
1²" (z, 2 -- Z s=),,
as when fixing Z-st, the exponential reweighting terms cancel out in the numerator and denominator the definition (E.34). Using the above identity, we are able to expand prt (z X X-s.) as
D (z X-ss) = Σv pT (z --s+)
x (Z-s (x
2-8
s. = X-s)) Σ [=| Z-sss) (Z + 1¹7 (X ( - 1). 1(2_s-
1+ 11 (x((X (X-s. (ea 1)
2 S
(2) + 1et
(x = s.)) 11 (X (e -1)
1 nTTH (X s.)
(ea - 1)
53

---

where the second equality follows from the fact that the reweighing term in z=' lifts the likelihood of Z-s- s- by a factor of e" relative the base distribution μ." (Z-s+) ). and the denominator is just the normalization constant. In the sequel, we let r (x= a) (1 + nm (X-s.) (c" - 11) be the inverse of the normalization constant. We then have
(=| = - 1e (=) + (1 (1-11 r" (X-s:; (x (x= = X-s)).
(B.35)
On the other hand, by definition of i=" in (E.34), we directly have
1. (x"(X-s') )aa = - 1 = Exn (z(zzss)) exp (a allness 1(2-n x-h))
(X 1º (X-s.).
(E.36)
Combining both (E.35) and (E.36) have for ss that
1x = Ex | X-s))
(1-r((xxx))) (= (x= = c/ X X-s)) -1) 2
Js E
rdt (X-sx;;30) a) (x= + -1 r-
kE|d
(z = ex X 77" (Z-s. X-s. (x-s))
e X (I == (x)
= B g mpp
p (x((X++::a) a+ 1 (x (x) + (1 pT (X a)) 1st (x= ck X-s.)
kEld
(T" ex X-s+) e0r (X (x=+);00² 0)) 11 (x((xs)))
(x= ek X X-s+)
1. (x =x))
-e (X 10 (X (x)s:) 1
E g+ mpp
(2 ck X-s)
kEld
5- a, T)
Here, we note that J(:; a, .T) is a function depending on both a and n, and can be expanded as
(n = ex X-s.) - 1= = =x))² r (X (x:)a= ex X s+) + r" (X s- (x) ex)
J(X Sk a, x)
kE(d]
Hence, we complete the proof of Proposition E.4.
Proof of Proposition E.5. Also, note that J(X-s.jr,z) is convex in r, as by taking the derivative of J(x-s+;r,n r,, т) with respect to r, we have
aJ(x sT; -);35s);,
6k X s.) x-)
11th (ex))3 X X-s+) ) + T/1 (Ex))2
dr
((1 - r) (x= =
kE|d
82 J(x --X---1 s.TT,
(n (x = ek X s.) n3 (ex) 2
X + rM (ex))3 2
dr2
((1 - (x
kEld]
Hence, a naive upper bound for ((_=s;;,,)) is
J(x ((X-s+;7,m)
S max ; 0, T), J J[X 1+
< max (D (1 (+) 1h (-1 (1I X )),
(1)x))))))))()))))
54

---

where we remind the readers that Dx=(ullz) = Ex (u(x) - ((x)))/v(xa)) Next, we show that T(--s-;r, π) is Lipschitz continuous in r:
(J(X-s+;7, dr
X s+) -²
kE[d] - (1-r) Eh X s.) + rpp (Ex))
(n- = E/ X so) p (ex))²
ex X --ss -M
r)n (x 6b X S*) + ryc (ex)
1-rr)y ek X-st) ) ryr
kE|d]
1 ex X-s.)
(ex)
- (X-s+, max
(ex)
(= = Eb X-s))
S max (Dx²(1 (-1x-(()))) Dx= (x pº (1))
where we use both the upper bound for J(x-s+;r, σ) and the lower bound for the transition kernel that both μ² p=(11 (-)) and μ²
p²(+)) (-) are bounded between y and 1.
E.5 Lemma on GIH Approximation Error
Now given the convergence result for the training dynamics, the natural question ask is how well the learned model implements the GIH mechanism. In the following part of this section, we state the lemma the approximation error and also present formal proof of the lemma.
Lemma E.6. Suppose Assumption 3.5 holds and consider training a transformer model
TF(M, H, d, D) with H = M. Let
A1:11-ps: (ti), 22: II (o
KES"
where t and 12 are the ending time for the first two stages of the training, respectively. Suppose the error A1, A2 OLI-1) after the first two stages' training, and O(lo L) after the last stage training. Let y be the output of the model (2.5) after the training and y the output of the GIH mechanism (I M, D) defined in Definition 3.2. Then for any π E supp(P) and with high probability - OLL ) it holds that
lly U111 O(L L) L.
Proof of Lemma E.6. Let 87 = I ( = I1+1-A) and 8t = (ui+1, u). Invoking Lemma F.1, the model misspecification error is bounded by
max M<ISL
|sj stl s (A1 + A=):=A. A.
(E.37)
We note that the second layer's 's attention weight a can be as large as (log /8. We are comparing the output of the model with the GIH mechanism GIH(#1:1: M, D). Let N Ei>m l 1 x1+1--). The output of this GIH mechanism is given by
N Σt Ei-M41 M+1 (T. M) Et= It, otherwise.
llnss 1(ai- #L+1-b), if N N21,
1/
We define
[ . Il ( #L+1-k), if N21, 1,
of (L M)-1, otherwise,
with o² = (af)s Since ||ل|||| = 1, the f-1 norm of the difference between y" and the model's
actual output is given by
lyy wlli S 110- - -o11
Let us define the set r = [L21>> M Ihes- 1(21-1 F1+1-1) and (12 M
Thes (xi-1 #1.+1-1) 03. Using (E.37), for 1 E r. we have 1 2 2 88 - and
for l E T, we have $ A = A. Consider the normalization factor in the softmax function.
Σ M=++1
exp(a 8).
55

---

By the split of the set G and I' and noting that [T] = N, the normalization factor is lower and upper bounded by
2 2 N exp(a (1 - A)) + (L. - M - N) 2-, Σ exp(a) + (I. - M N) exp(a A) =: 2+. Let us consider the event N 2 in the following. We then have for that
exp(a st)
1 -
exp(a)
1 V
exp(a
(1 (-))
loo 0:1
Z
2
Z+
1
N Nexpl -AA + (L. M -N) exp( -a)
N
V
exp(-2aA)
1
N Wexp(-aA) + (L - M N) exp(-a)
N
N (1 - exp(- + (L- M - N) exp((-)
- exp(-oA)
L L. exp(
(N exp((-aA) + (L M - N) exp(-a)) N
N
N2 exp(ppaa>) aa
Note that aA = o(1) due to the assumption that A (1- and (L). The right hand side
upper bounded by O(a (N) e(Lxxp(a))/N-) For /
we have
exp(aA) 2_ <
exp(a (2A 1))
exp(-a) N
of - 01 = 05
In summary.
Ily" - ~ll1 S -o11 - Σ
|oj - oil m/++ +
IET
LET
N (ANN Lexp(-a) LAXN-L)x((o))
(exp(-a)
5o + Lexp(
+ 1.
N
(E.38)
The above inequality holds whenever N 2 Now we aim to upper bound the probability that N = 0. Note that N Et 1(XI-s- X1+1-s+)) We consider the following second moment:
E {[ (L - M) Σ
1(Xi-s = E) -(=
] < Dx² (L (L-
Σ
1(Xiss = (00()((())
I M-1
I M=++
M
WE E 21551
L(1 - x) 1//-1/22
where the first inequality holds by noting that Dy= v) (p(x) - (2x))² and the last inequality holds by Lemma F.18. Therefore, by the Chebyshev inequality, we have
P
L
Σ 1-1
1(xiss: =E (E) - 2t)
S L(1 A) 7 8-1
We can take / = minisexis" 1= (E)/2 and by also taking a union bound over E X16" (which gives a s factor), we conclude that with high probability Ø(1 - L-1) it holds that N 2 tL L min Fex1s-1 μM (E)/2. Thus, it follows from (E.38) that with high probability
lly* vll, 0 as +
exp(-a) min 1 (E)/2 O 0( (L-1log I. + L L-a/llal///// log
Hence, we complete the proof of Lemma E.6.
56

---

Auxiliary Lemmas and Their Proofs
In this appendix, present the auxiliary lemmas used to derive the approximation of the gradient flow dynamics in the proof of Theorem 3.6, which presented in the previous appendix. The proofs of these lemmas are presented right below their statements.
Useful Inequalities
The following lemma provides bound on the model misspecification error, which is the difference between the model's output and the ideal output.
Lemma F.1 (Model Misspecification Error). Let ui+1 be the output feature after the FFN & Normalization layer. Then, the model misspecificatioo error defined as
max [E(L]
(uL+1, ui)
-II 100--=II41--0 hES*
bounded by A + 4 where A and 22 are the ervors after the training of the first and second stages, respectively, and defined respectively as
A - p
--p''
12:=1-
Il
(a )
bES
Proof of Lemma F.I. By definition of the output feature un after the FFN & Normalization layer:
(uL+1,2)
Σ Ps
II (h)
Lo)
2 (11)).
SE
AES
As each w/a) is a convex combination of XM() where M(l) = (M-M,MM.... M,..., 17. 1|0 (2)
1l S 1.
Thus,
(UL+1, un)
II BES
(u)
(h)
=
Σ
Ps
(0)
(h)
II RES-
(0) (A) -2(
SE(H)so
RES
x
(1 - ps-))
(
(241)
(h)
Σ Se(H]gn (S8*
Ps
II
(o)
(h) (h)
hES-
hES
S max
1 - Ps
Σ
Ps
1 - Ps. =
SE
where A is the error after the training of the first stage. Since (h)
EjENs o x1-jj we have
II BES*
(1) (M)
II
Σ
0-0-- (1--1)
(h) (h)
KES
(JE[M]
Σ
II
(b) (h) o 1(x1-11 = #1+1-5h).
(inss))esss EM | KES.
Here in the second equality, we exchange the order of summation and product. The last term of the second equality can be understood as follows. We first pick (S*) index pairs {(hn, arbitrarily,
with each 7 ja [H] Then evaluate the product Ihess- - o o- 1 (x2-i, x1+1-jj.) given these indices. Then we sum over all possible values that jih, hEs take.
(h) (h)
The above equation implies that
II (u)
(K)
(b) ULt
II BES-
(h) (a-
(a
1 (0L =
kES
Σ (Wndhhees ((h,)))ees hES.
II
σ --n J-jn 1 1(11-in = #1+1-51)
(h) (h)
<
Σ
II
(h) (b) o th
<1 - II
(o =:
(F.I)
(ind)))eer ((h.)))e KEST
hES
57

---

where the last inequality follows from the fact that
Σ
II
II KES"
Σ
= II
Σ (E[M]
o
O
=1
(BadJa JAES" BES'
AES
Here the summation sign in the right-hand side of the second equality indicates that in the last line of (F.1) we sum over all possible values that Ji) hies- can take, except for the only case where (bh. Jh) = h) for all h E [H].
In summary, by triangle inequality, we have shown that
(UL+I.
1(x1- = zL+1-N) SAi+ A2. II
hES
The proof is completed.
Next, Lemma F.2, we establish a uniform bound for the quantity involved the gradient.
Lemma F22 Let y(k) Et- or 1 ex) for each & [d] where E1=M+1 of and ai 2 for all [L]. Let e and C be two positive numbers. For any C-hounded function X [-C, C}, we have
Σ
L
o
1 = -I -
((X) y(k) ((T) [EW)
S2C.
(())+8
((k)+ +
1=M+1
Proof of Lemma F.2. By the triangular inequality, we have
(x1+1 =x)
y(k) 11 (x = =ec)
F(X)
(kk)+ +E
y(k) +E
M=
k=1
1(x1+1 = ex)
y(k) 1 = cx)
C.
1(x = =ex)
+
Σ ai
y(k) +
u(k) +e
k=1
M=++1
k=1
y(k)
1x =ex)
2C
$2C,
(kk) +
where the equality, we use the definition y(k) = E o1 = ex) and E/= 01 Now we conclude the proof of this lemma.
F.2 Approximation Errors for Dynamics Analysis
Next, Lemma F.3 addresses the approximation error induced by di = 1/L the transformer model. The approximation error will be for go.s to gi.s for Stage I and Gh.1 to 9h.2 for Stage II.
Lemma F.3. For the transformer model defined in (2.5) and any bounded function x R
such that supzex [f(x)) S C for constant C O, define two quantities and B as
[ 21 =ck
y(k) (FL+1 1(X),
A
ExiT (i(ss)
y(k)+ +
u(k) +
L M++1
kEld
1
1 1t = =x)
7(k) 1 1x = =cx)
-[x]
B
ExiT
M
8(k) +
v(k)+ +
I=M+1
kEld
where s = TT UT and y (L - (L-M)- E1=044 Tt. Then, for all E [0, 1] and e E (0, 1]. it
holds that
8Cad
A B| <
58

---

Proof of Lemma F.3. By triangular inequality, we have
L
1 3j ck.)
A
S
E
(a 8) - I-- M
ukk) +
1= M
+1
kE(d]
1
[ =1 = ====0 =kk)
[x1+11 = 37 ck)
M
y(k) +E
7((k) +
(u(k) [ =ck)
+
(a 5)
L - M
2(k) +
1
y(k) [ = ex)
T(k)
[x ck)
}
s(x)]
M
ukk +
v(k) +
Note that 0 < S for all I = M + 1, 1,... L thanks to the layer normalization. Then, for the softmax operation, have
exp(a)
Soi(a.8) S
I+ (L - M 1) exp(a)
L 1 cxp(a'
which implies that
1
cxp(a)
(a s)
S max
I.
M
I. M
1+ 1 - M 1) exp(aa)
T- - exp(a)
I. M
exp(a) -1
(F.2)
L M -
Since indicator functions are bounded above by 1, we have
11 [1 cx)
1
u(k) (xL+1 = ex)
1
(F3)
v(k) +
u(k) +
For the second term, we have
I (ZL+1 = ex)
1(11+1 = = 6x)
[7(k) -y(k)]
(as+ilo((a-s5) - (L - M)
(lk) +
o(k) +E
22
g²
exp(a) (a) -1
(F.4)
e2
where the last inequality follows from (F.2). Similarly, the following bound can be derived:
y(k) 11 (xI+11 = ex) v(k) 1 (x1.11 = ex)
exp(a) -
(F.5)
ukk) +
5(k) +
Combining (F.2), (F.3), (F.4) and (F.5), it holds that
A - B
<
Σ
1
E
exp(a) 1 (2(T. - M)
s(x) <
4Ci(eep(a)) - 8Cad
E² <
g²
M=++1
kE(4
where the last inequality follows from exp(z) 1 < 2x for 0 x x S 1. This concludes the proof of the lemma.
Lemma F.4 provides the approximation error introduced by (ex) R T(k) in the transformer model.
Lemma F.4. For the transformer model defined in (2.5) and any bounded function f : x + R
such that supzExz [s(x) S C for constant C O, define two quantities A and B as
1
Σ
L
Ex|-
1 = 21 = ek)
T(k) 1 =x))
F(X)
L M
v(k) +
u(k)+ +
t= M+1
1
[ = = ex) -1) s(x)]
B =
Exisr
L M
nT (T+(e)))
L=
ke[d
59

---

where y = (L MM Ei= 21. Under Assumption 3.5. it holds that
|A- BI S 4C (1 (D (r)) + 1)1/4 + 2/M + Cy
1,1/2
where ( is the initial distribution over the first rn tokens X1rr. Here we let Dxa (Po ll p²) to
denote D D (No(r = I1 (X =)), i.e., the x -iviiergee between Pn and the distribution over the first Tn tokens under the stationary distribution p².
Proof of Lemma F.4. Let us use yx (-) to remind the readers that 7() is also a function of X. We simplify the expectation [xx1== by E in this proof. By rearranging the terms, we have
L
(xL+1 21 Ck)
( x1 = ck)
A
B
E
Σ
L - M
yx (k( +E
1 (=(())
Im M
+1
AEEdd]
kE|d
(x(K) [ =(x ek) +1)
-(x)) -
Σ
Ox (x( +E
kE(d)
I.
1= (ex) - (x()))
L
M
E
(ūx(k) + 8) 10tt (Ex)
(kx + +) 1 (Ex))
1= MM
+1
kE(d)]
e 1 ek)
-fx)
(x(k) +
kE(d]
Here, we have three terms to control. For the first error term, we define
1
1 (ex) - 5x (k)
[0I+1 = = =I s(x)]
erri
L M
E
[Tx (k) + +) 1M (ex)
i=M+1
kE(d
C
(ex))))) xx (k)
E
1(1+1= = x1 =x
L M
(Tx (k) + 10a (ex)
M=
kE(d
11th (ex)
Tx (1:)
S C. E
1(x1111 = ex)
10 (Ex)
kEId
The first inequality above holds by noting that supx [x(X) S C and the last inequality holds by (L-M)
noting that (x(ee) = -L-Mi Et a0+1 1(x Ex). Using Cauchy-Schwarz inequality, we arrive at
27 1 (ex) (x())
1 cxx]
(ex) =ca)) 1)
1/2
err1 C
E
(ur(ex)
kE(d)
kE(d]
S Cy-1122 1/2
E [xx (Ux(+) 11 (21+1=
For the second term, we similarly have
E (ix (k) + (ex)
1(11+1= = =1 =x)))) s(x) =(x))((x)
err
E
L - M
M+++1
kE(d]
11 1xI ex)
< C
E
SC-
vT (ex)
kE
Lastly, we have the error term
erry _- M
1
Σ
E
1(1141 =c)
-(X)
5 SC
el c = ex)
(x(k) +e
Tx (k) +
1441
kE(d]
kEid
1 (XI+1 Ex)
Te( (k) - 10 (Ex)) ( ex)
C.
E
+
Σ IV
pe +e
(ex) +
kEId
kE|d]
60

---

Here, the first term is upper bounded by Cy E, and for the second term we have by Cauchy-Schward that
(yx (K) nr (ex)) 1x1+11 ex)
=ex)
C
Σ RE(d
E
(ex) ) + ( (k) +
Ux (l)) 1= (ex))
E2 ( = c+)
SC.
E
(ex)
(Tx (k) + (cx
AE(d]
kE|d
S Cy 1/2 Ex xxxxxxxxX) ll (1 (#1+1=
which shares similar upper bound as erri. Now we invoke Lemma F.18 8 to conclude that
A - B| emr + CIT2 + erry I 20y-1/2 1/2
/
pe (IL+1 = +
20y -1/2 4(1 - 4(1-A)-1 VDza(Hu (T)) +1+ 16M
1/2
+Cy
L min 10t (II++)
20y 2(1 - X)- (No 11ah 11 + +4VM
+c-1le.
1,1/2
Hence, we complete our proof of Lemma F.4.
Lemma F.5 covers the approximation error due to the mixing property of the Markov chain. Lemma F.5. Let S E [ be a fixed set. For any h E S, let ((h)) and (h) be two fixed probability
distributions over [M]. That is. for any i.j E [M], we have ((h) ; 0- E [0, and Ei=1 1010011
(h)
(N)
EN 0 = 1. Given these distributions over [M]. we define
(b) IL+1-6: and SE[M
Σ
1/
Σ JE[M]
where let X denote the 1-th token in the Markov chain for all E [L+1]. Moreover with slight
abuse of notation, we let (2, Z) = (z, z 1....... 2_M) E X XM+ and (x, (x, (I...-1..... .-M E xM+1 be two independent random variables sampled from the stationary distribution v², We define random variables ((k)(2) and w(h) (X) as
((N = Σ
(h)
2-i.
and
r xx) = Σ
M
SE(M]
Using 1 oj
(1h) (N)
g0h (Z), and (hh) (X), we define two quantities A and B as
A = L M Σ ExIT
1 = x1 = ex)
-1)
II hES
(ut (h) ((h)
1=(ex)
I=M+1
kE|d]
B = (e xp
1(x (=====0 (====e=)-)
II RES
(7(h) (Z), (00) ((1)
1rt (cx)
kEld
where Ex|a means that the expectation is taken with respect to the randomness of the Markov chain with transition n. Then, under Assumption 3.5, we have
8M 16 (GDDx= 1 L
A-
L(1 1811111
where po(-) ) is the initial distribution over the first ra tokens Xirr, and D 11 1²7) is a short-hand notation of Dx2 (40(Xir.. = (X1.r. +)).
61

---

Proof of Lemma F.5. By triangular inequality, we have
11 (xL+1 = =x
A B <
Ex|T
II RES
L
M
(T(())
kEId
1(x 1=====eN)
(II hES
w((h (X)))
(e.,
((E))
kEId
Σx
II TES
(o) (N) (1)
(I (6th)) (2), w(hh (x))
M
(=M+1
We will establish the upper bounds for each of the absolute value terms. We first focus on the first absolute value term.
Bounding the First Absolute Value Term. Let (X) denote the joint distribution of the whole sequence under kernel æ. By the definitions of 2/ and of we have
(h) (0) L-11
Σ
(h)
(h)
-jh
(I-+1-
Σ
- 1(ILtI-i = Ii-ja = ck),
J[M kE(d
where we use (hk ) as the indices to highlight that they are associated with head h. And we use kn E [d] to index all the possible common values for Ti-á, and FL+1-jk- Then plugging this equality into
(1) and exchanging the order of product and summation, we have
1(x1-1 = -I - =x==
II hES
(h)
(F.6)
kE(M
Σ
Σ
1(I++1=20) = =I
(h) (N) 1(IL+1-1 = I1-ja cka.) II hES
(= EN)
[GnnJ)hess (Hel
where the summation means that we sum over all possible values that {in, Jn. kb Jagn and k can take. Specifically, each in and in take values in [M], and each k,, and & takes values in [d]. Moreover, using the property of indicator functions, we can further simplify (F.6) by gathering all indicators:
(x1-1= = = ex)
II hES
2t11
(F.7)
(F(())
kE(d
Σ
II bES
Σ (kE]
1 = Cky FL+1-is = I1-5n = Ck. vh E S)
-in
11 = Ck)
Now we take expectations with respect to the randomness of X on both ends of (F.7) and get
Σ E
( =m=
II hES
(u) (b) (h)
L M
(xx)
AE[
El P. (EL+1 = F1 Ck, FLt1- = 31-5 = = vh E E
Σ
II (h) (h)
Σ (habeeskeld
-Jh
(L- 1= 2 ex)
((bsJa)eee AhES
To further simplify the above equality, we define a new probability distribution over X M:L++ and another subsequence of length M Note that X M:L+1 contains is a subsequence with M + 1 tokens. We let (z, 2) = (z, 2-,,........ 2-1, z-M) denote a random token sequence of size M + in reverse order. We define joint distribution p over X X McL+1 and (z, Z) as follows. Let E (Eo, E-1.... E_M) and E" (E,, E" 1 ...... E" .)) be two elements in M That is
62

---

each component of E and E' are in A. The probability mass function of pa is defined as
Trt ((FL+1, FL. FL..... FL. FLLI (2, E")
(F.8)
I
Σ M=++1
pa (
E, M) m1- (.1.. E)).
L- M
That is, f can be viewed the joint distribution of -++1-M-L+1 with averaged distribution of the history. When L is sufficiently large, by the mixing property of the Markov chain, we expect that, under pr, (z, 2) is approximately independent of #L+1-M-L4H and the marginal distributions of (z, 2) and XL+1-MS M:L+1 are both close the stationary distribution μ², We will translate this intuition into rigorous argument in Lemma F.17, which bounds the total-variation distance between pn and the product distribution jº p=.
With fn defined in (F.8), we can rewrite the expectation above
Σ
E
1 =I
II hES
(h) (o) . ..411
(F.9)
L
M
(=x)
M=++1
Σ
II
(h)
Σ
(01+1
= =
Ck: FL++1 = =z-in Ck. 4h E S)
-in o
-jh
1 (z = ex)
Similarly, by the definitions of 7iA) (Z) and (1() (Z), we can write (B( (2), x(N((x() as
( (2), v (X)) Σ Σ
aiin Jjn 1z Ij.. ex).
(h)
(h)
[ E,E(d
Then, multiplying these terms with h E S, we can write
1 (x = == ck)
II (7th) (Z), (())
(Z),
(=x)
keld]
hES
Σ
II bES
a (b) -sn 0 --in)
Σ kkil
1( x Cb. 2-sh = =Z- = ebo 4h ES)
10 ck)
((inWhhes
(F.10)
Recall that here (z, 2) (z, 2-M.....M-M and (x, (x, (x. x x-M are independently sampled from the stationary distribution p². Taking the expectation under n², have
1(x 1 ==
an))
II (T(h) (Z), w00 (() ) hES
(xx,)) Z) x
(F.11)
1 (cx)
kEld
= =
II ((h WEs VRES
(1)
Σ
1 Ck. 2-14 = Ckni vh E () (1 Eks 2-th = Cka vh E S)
-in
-Jjn
(= = Cx)
To bound the first absolute value term in the upper bound on IA B). we aim to compare (F.9) and (F.11). To this end, let us fix collections of index pairs (in hes. Let S {in h S} and (in)RES and ( Jhss Since there might exists two
S = (in h S} be the unique values elements h and h in S such that in = in or jr = Jur, [S1] and [Sy] might be strictly less than [S]. As result, fzd (x L1 = Z Ck, UL+1-in 2-jn Cky Vh S) only involves random variables x # X1+1-S, = { 2, Z (2= which are subset of the random variables
defined in (F.8). Similarly,
(x Ck, I-I EKA, vh E S): 11" (z ek, Cka Vh S)
only involves subset of random variables X-s, = (x-ities, and Z-s.- Let define E (Eo, (E_i)ess) E S/1t+1 and E' (E,, (E',)jess) x|S1+1 By enumerating E
63

---

in X/ and E" in X1 we equivalently enumerate all possible values the above random variables can take. Therefore, by comparing (F.9) with (F.11), have
Σ
(=
Cb. FL+I-in jn = Cka. wh E S)
(x = ex, = Ck,, vh E E 1==
Cky 2- Ck,, 44 ES)
Σ E
1e (141-)) E, (.,) = E²) - 1 ([IL+)) = E) 1 (Z, =E))
1(6 Eg: E-i = jn Yh E S) -" -")
(F.12)
where in the last line, we use Y and Y' as placeholders for the random variables (IL+1, s and (z, Z-s,) respectively. In the first equality, we sum over E E X|Si++ and E' X| and the last inequality follows from the definition of total variation distance and dropping the indicator. By Lemma F.17, this total variation distance is bounded by
2 2 (Y Y = - 1 (Y = x 10 (Y- = l|TV
4M
8 DD (10 11 +1
L(1-A)
min X (FLL1,
4M 8V BDDz2(10) 1=)+ +1
$
(F.13)
L(1-A) 7/1 +1)/2
where the last inequality holds by Corollary F.15 and the fact that |S [8). Specifically, Corol- lary F.15 implies that the density function of the joint distribution of 2141 and X is lower bounded by 7. 111++1 2 7 SI+1 Thus, combining (F.12) and (F.13,, we have
Ek: IL+1 1-ih i-jn Ck,, Vh S)
Σ
1101 = = 6%)
(kn, (Hn
1 ( Ck. Ck.
Eks wh S)..p(( Ck, Cba vh S)
1= = Ex)
(bnhnes.ke kE(d]
< 1
11t) ) 7 I L(1-)) 1(1)++1)/2
4M (D
(F.14)
Therefore, to bound the first absolute value term, we combine (F.9), (F.11), and (F.14) and use triangle inequality to get
= =1 -
ExIT
II WES
(m) (1)
L M
100 (ex)
L= M+1
kE
[E kE(d
1(x 1(x= 1(==== ew)
((n,,)) X
(Z), (1) (X)))
(ex)
hES
(.(4M 4M 8VDx ll (1T) +1 < 7 +
1
Σ Hiies
II hES
o (h) (h) in -ih
(F.15)
L(1 A) -
Furthermore, recall that ((k) and (h) are probability distributions over [M] for all h E S. By going over all possible values that {(hn. f1)) Tnes can take, we have
Σ (hEE
(II (k) (N)
(
II KES
Σ -
(
(h)
1.
(F.16)
RES
kE[M]
kE[M]
Plugging this equality into (F.15), we show that the upper bound in (F.15) can be reduced to the right-hand side of (F.14).
64

---

Bounding the Second Absolute Value Term. For the second absolute value term, an analogous argument can be applied. In fact, the proof is simpler because we only need to handle (o) (h) (h)
and (Z), (l))))))) (2)) and do not have indicators Il (xL+1 = m1) and I (x = 2).
Similar to the derivation in (F.9) and (F.11,
Σ
E
II -EES
(N)
((h))
- E(r,))) Xp"
II
(Tth)) (2), 21th (X3)
(2),
L --
M
t=
M1
hES
Σ
II hES
-in in
(1) (b)
a -Jn
Σ
(2L+111t = =2-h ekn vh ES)
(F.17)
(HADNES
- Cka. vh E ES 1 (z-jn Ckn wh 5))
Similar to (F.12), for any fixed collection of index pairs (in. jnJNEs, we let S₁ = fip h E S) and
S² = (in : h E { denote the unique values (ix MES and (jn) By Lemma F.17, we have
Σ 0 (1 = 2- = =kk wh E ES - 1 (x-in = Eku wh E ES) 1 (z- = Ekpp vh E ES)
(Anlies
4M 8/Dx2 (,T)) +1 S
s <2ppp - =
= - (Y'
L + L(1 A) 7181/2
(F.18)
Here we use Y and Y as placeholders for random variables X and Z- We note that Lemma F.17 can be applied to any subsets of X XL+1-M:L++ and (z, Z). Therefore, combining (F.16), (F.17), and (F.18), we conclude that
1
Σ
E
II
(N) (o)
( (Z), 0(0) (()) (X)) )]
L-
M
I=++1
MES
hES
4M
8/
8
D22 1T) I
L(1
Xy/55122
Note that the second upper bound is dominated by the previous one. This completes the proof of Lemma F.5.
Lemma F.6 provides an approximation result using the definition of the modified 2mmutual informa- tion.
Lemma F.6. Consider a fixed set S E [H)DD. For any h E S, let (()) and ((h) be two prob- ability distributions over [M]. That is, for any i,j E [M], we have δ ((h) (h)
--:00 E [0.1], 1). and 1. Moreover, we let (z (z, (2,2-M
3 =E (1)
M (k)
M+1 and
(x, M, x-1) xM++ be two independent random variables sumpled from the stationary distribution pta We define random variables g(h) (7) and wy (X) as
7(hh) (Z) = Σ >= 2-1; and
(x) Σ c xz--
xz---
iE[M]
JE[M]
Let (it,, (h,.i Nes be any fixed collection of index pairs, where it, [M] and jī E [M] for all h e S.
We define quantities A and B
1(x 1k2= ====ei Ex)
-1)
(1)h)) (Z), (2) (x))))
(x. (x) xp²
(Tx)
KES
1(x =2 z= Ex) -1) 1)} (Ex)
B = ET, (2.)
1 = =2-5;)
heS
Under Assumption 3.5. it holds that
Ex [A]
II hES
(k)
(h) 33
B
<
(1
- -;; -x7(8*)
(h)
hes
65

---

Proof of Lemma F.6. To simplify the notation, we define a signal set ((S) and an error set ((S) as
T(S)= == (5ns))
T(S) = [in. (in, E (( x [M]] \T(S).
Similar to (F.10), we can write Ex [A] as
B, [A] (x,, xp²
Σ
Il
(h) -in -Jn 1(I-in = =Z-sa)
E kE(d]
1(x= =
-1) 1)
1 (ex)
(onddahhee KES
where we exchange the order of product and summation. Using the notation ((S), we can split the summation into two parts:
E [A] = (1x) Xn
II hES
(b) (h) - o 5h
1(x-;;
= =2-5))
1(x = ex)
- 11
1e (ex(
b [c)
1(x 1(x= ======a) ex) -1) ]
+ (ox)) xp
Σ ((b.)a)))e hES
II
(h) (b)) 0
1 (x--in = z-jnn
Σ
-Jh
kE(d]
(h) (h) II hES
B
% -5h
+ (a, (ox
Σ
(b) g(0) II
Th
-ha 1(r-in =2-j))
1(x Ex)
-1) :)}
(Tx(e))
(( hES
kE(d
Here last equality holds by the definition of the B and the fact that (A) and (hh) are fixed vectors.
Therefore, to prove this lemma, it suffices upper bound the second term above. To this end, apply Lemma F.7 stated below for any fixed set of indices (in, jn JaeS E T(S). Specifically, let Si (ix: h E S} and S2 = [jn: h {} denote the unique values of (in)tes and (ja)NEs.
Lemma E.7 implies that
(1 x
II 1(r-n hES
E kEld
1 ck)
-
s (xa(S)). (F.19)
(= = cx)
Combining (F.19) with the fact that
Σ
II
(h) (h) o -in
II hES
(h) (A) -i7 -j
=
(Hn.A)RRCSE( hES
the desired term is bounded above by (1- (1-IInes o
-n -35. Ixx (S*), which concludes the proof.
(h)
Lemma F.7. Let S E &D be a fixed subset and let {(i,, jn Jives be fixed collection of index pairs, where in. E [M] for all is E S. Let S = (in: h E S} and S₂ = {iu: h E S) denote the unique values of (in)ee and (jk njs) We let (z, = (2.2-M....,,2J)
XM+ and
(a. X) (I., -I, .........22 x-1)) M+1 be two independent random variables sampled from the stationary distribution p², where is the transition kernel of the Markov chain and is sampled from prior P. If Assumption 3.5 holds, it follows that
1xx 1x==== =ex)= -
-1)]
EgpP Fe.P((,,))tee (1) x
II hES
Σ kE(d]
(= ex)
S 1
+ + +Iz((5)) S sTx(S*);
where (S)(S) is the modified x2 -mmuuaal information defined in Definition 3.1 and St argmax Tx²(S).
Proof of Lemma F77. We first note that it is allowed [S| * |82] as there could be duplicate values in both (ix) and (jn) hess while Sj and Sn are the unique values. In the sequel, we let X-s, denote {i- his and let Z-s, denote {2-5n hess where repeated elements are removed. Moreover, we let
66

---

(X-s, = -Z-s,} } be the event that i-in 2- for all h E S. Notice that I (x =2-0a
1(x-s = 5.) Then, we have
1 - ex)
-1) 1) ex)) [1]
II -ES
1 [ (x-in = =z-10)-
(= (====e)) cx)
(F:20)
REId
1(x
-=,,),.)),,,,)),),, 11 (xx-s =(-sa)-
15 = ex)
=ex)
kE(d)
1= (x= = exl x (2 = en cx|Z-sa) -1) 1(X-s, =Z-sa)
Σ kEld
10 (z = 64)
(= = ek 2 (xs) -
1 (2 ( (1"(2 = ek ex|2-8a) -1) 11" (== =ex) 1(x = 22..))
= -x) XM²
11a ck)
ck)
kE(d
Here in the second equality, take conditional expectation given X X-s, and S- The last equality can be verified by direct computation. To simplify the expectation above, we aim to transform the indicator of (1 Z-s.) into probabilities involving X-s, and Z-s,. To this end, need to explicitly enumerate all possible values that X-s; and Z-s; can take. This challenging, as there may be duplicated values in both in and jn. and thus X-s, and z-s, can have different sizes. However, since S is a "reduction" of (ise we can revert to the original space and consider E = hER)NeS E 215 that respects the reduction from hes to S1. Here each E,, is the value Ii takes. In other words, with (én Jes that might have duplicated values, we consider the values taken by (z;, Jas, with duplicates allowed. And E has the same duplication structure (2t, hes. In the following, we describe these values by introducing the notion of compatibility.
Definition F.8 (Compatible Value Set). We say that X compatible with (Ex)EeS if. for any h + h' such that in = ig, we have En = Eng In other words, the unique values in É can be indexed by {in Jhes Si if E is compatible with hee)...
By this definition, E is compatible with (ix)tss it respect duplication pattern of (Ii,) Jes. If in iv, then we know that 2i, and Eig, is the same token. Since 2ú, and Tiy, take values E₁ and Eve, we must have E,, = Eg.. As a concrete example, suppose S [1,2,3), and the values of (in)kes given by (i1, 1g)
have S₁ [1,2]. which contains the (1, 2, 1). Therefore, unique values of (i,, 12, 43). Now, let = (E,, Eg, Ez). For be compatible with (in) bhess we must have E E3 since i iz. There is no restriction on Ez. So, compatible value set for this example could be E (a, b,a),, a), where a and b are elements of X.
In the sequel, define E the set of vectors in XSS) that are compatible with both (irhes and {in Jhes, i.e.,
8 = F. ( E is compatible with both (in)es and (in)ees
The compatibility condition allows us to assign T T-i,, ; the value E, for all h E S when F. E E. Under this assignment, the constraint X-s, 7_s, is automatically satisfied. We use the notation {X-s, = E} to denote the event that x: v-1, = En for all E S, and similarly for Z-s, = E. In particular, are able rewrite the indicator 1 1(X-s; Z_s,) as ΣEEE 1(X-s, E, E,
67

---

E). Then we can rewrite (F.20) by separating s,s, and Z-s, as
1(r = ex) -1) 1) -1)
Sy,,,,,)(,;,)),) XM*
II 1x in. = ==)
Σ
(x(2=e() ek)
4=
-ES
ke(d
1t (= = Cxl S = =]
-1)-
(1=(== es =E) -1)
= =Eg Σ BEE kE|d
1 "x exk
(= eba)
# (x-xs. = =E aZ = =E) 10 cx)
p= Ek (s-s))
-1)
2
#1 = ex) (X-s. = Ei)
ENT
2
(x = t4)
EEE' kE(d]]
1 + ス Ex
(== = en (-s)
-1) (0 = ex)
(z ex) (Z-s, = m))² (F21)
EEEE kE[d]
where in the last inequality, we apply ab S a² + 12/12. Next, for each E' E E, consider E" = (Ei)iEs, such that E = En, vh E S.
(F.22)
Note that for each E E E, E' must exist and is unique. The existence follows from the compatibility definition, which allows us to index all the unique values in E by restricting the indices to the set S₁. The uniqueness is due to the fact that (F:22) completely determines all the values in E' because enumerating over in for h S is just the same as enumerating over i for i E S1. In fact, E' contains all the unique values of E. In the above example, we have S₁ = {1, and thus E' = (a, b) when E = b)).
Since E' is uniquely defined based E, we are able to define operator J, that maps E E E to E" E x11: according to the mapping given in (F.22). Let Ji (E) be the image of E under Jr. It important to note that for each
E (E), there also unique pre-image F. E such that (E)) = E" according to the rule (F.22). Therefore, Tn is an one-to-one mapping from & to (8). In the following, for any E" E (E), we denote by {X-s. = E'} the event where x= = E" for all i E Sj. Equivalently, we have I-in, = E En for all h E S. Thus, the event (X-s, = =E} is exactly the same as {X-s, = }} introduced above. Therefore, the first term on the right hand side of (F.21) can be reformulated as
1 (x= -
-1)
1)= 2 (== Ex) 1 (X-(S- = E
Ex
ΣD
10 (x=)=))) ex)
2
BEE kE[d]
E
Σ
Σ
1e = ex ex
11 (= ex)) (x=(X-s: =E')²
2
(x = ek)
('E
< 2
E
Σ
Σ
p= ex ex(x_s)) -1) 1) 1 1= = Ex) (n = E')8²| =E')) = 2 xx (S)),
1
(= = ex)
F'XII RE[A]
where the equality follows from the bijection between & and T. (E), and the last inequality holds by noting that (E) S XlSil, The last equality follows from the definition of the modified mutual information. The argument for the second term on the right hand side of (F.21) is similar, and we hence conclude that
II 1x = =2-m)) hES
1(== ex) -
1)]
< I
1 Ie (Si) + 2 (s.).
(z,, X
(= ex)
REId
Lastly, note that T,2(S) < ((*))) for any S E [Hig by the optimality of S* Hence, we complete the proof of Lemma F.7.
Lemma F.9 quantifies the approximation error from or of and y(k) y" (k) for Stage III.
68

---

Lemma F.9. For the transformer model defined in (2.5), define two quantities fi and fz as
====
E
Σ
L
01
[xL+1 21 =x)
y(k) 1(IL+1 = =
II hES
(vf (h) (h)
v(k) +E
u(k) +E
M+1
1=1
1( = =x= yr (k))) +
y" (k) (25+1 y*(k) +E
f2:=E
II ( =x hES
LI=++1
k=1
where the expectation is taken over all the randommess in the data, and
exp (a N =(1--==2-11-) ZL-1-1))
oi Σt= exp (a Ilnes (zv-h =x1+1hh)
V²())) Σ
cb))
1=M+1
with S* is the optimal information set. Under Assumption 3.5, it holds that If f2l - 12 (1 a) (A1 + A2).
where A1 := Ps- and A₂ = 1 Ihee- (o (
Proof. We separate the approximation error into three parts [fi f2l err, + err2 + err3, which are explained in detail as follows.
The First Error Term. Here, the first error erni captures the error of replacing l (01-K= T1+1-n) with Ihes (o(h), (1411 in f2:
yf (k) 11 [x = =ew)
err ==
E
(k)))) +
+E
(k) +E
1=M+1 +1
kE(d
(o) (h) (h)
II
1(a = x1+1-1))
HESt
KEST
Using Lemma F.2, we have
1 $2.
Σ
8" (k)k)
(k) + {
1=M+1
kE[d
Then using Lemma F.1, we conclude that
erri < sup
II (u) Lt1
(h) (h)
II (---==0111 < A2).
tE[L] kes*
hES
The Second Error Term.
The second error term characterizes the difference in az and af:
err = Σ
(op -1).
Σ RE(d]
= -
('( ex)
II (u)
y" (k) +
1 (k))+
+
1=M-1
hEST
To characterize such an error, we invoke equation (53) in Lemma 5.1 Chen et al. (2022). This lemma states that for and o* being the output of the softmax function with scaling parameters a for s and respectively, i.e..
exp(as)
exp(as*)
Σ exp(as:)' and
Σ exp(ast)
holds that ||o ---0" -01,, S 4a
||s s* ol.o. Consequently, we have l1o - - S 4a (Ai + A2) by
Lemma F.1. We notice that
[(IL+1 21 ===== ex)
(k) (0L-1=EE) (k) +E
II hES
(o) (N) (h) maxxfe 1,11 = 2741
kE[d (*))) +E
Thus, err2 < 4a (Ai + A2).
69

---

The Third Error Term.
The last error term characterizes the difference between y" and y:
err3 =
Σ or
E
(k))+ (T-K))+
y(k)+ b(k)+e)
1(1++1==1 =x= =x
II hES
(o) (b) (N)
M==++1
ke(d]
(*(k) +E y(k) +5) 113
(k(k)
y(k)
Il KES
(v)
(h) (h)
M+1
kE|
By noting that y* = Ei=
of and y Ei=M+1 0101, we have
4
llv" vll -
Σ
)m
S
Σ L=
loo - 01l2 11x1l11 S -10000011 s 4a (41 + A2).
1= 1+
The first term of errz can be bounded by
E
Σ O5
Σ
y(k) +e) 111111 x1 :
II NES-
(u) (4) (h)
"L+1
1/" (*)
M=++1
kE(d]
lu(k)
(k)]
(k) + (vk( +)) v(k) [x ex) S 111 y² || . S 4ae - (A, A2).
RE(d
Moreover, for the second term of erty, have
L
(k)
v(A) v(k)
E
Σ
3/" (k) +e u(k) + 1(xL+1= =6) II (u (h) (0/1))
1=
M+1+1
kE[d]
hEST
ly(k) (k:)
1..11(2ag) 1 =xx S 4aE - (Ai + A2).
(vr(k) + (bk +
kE[d]
then holds that
If1 - f2l S err + CTT- + erry s 2 + A2) + 4ae + A2) + 8aE Sas + A2)
= 12 (1 + as= )) (AI + A2).
Therefore, we complete the proof of Lemma F.9.
Lemma F.10. Let us define for brevity,
11 Z) exp (a I 1(2 = x1+1-1)) (2 =)= Ftt (2,2(((41-ss)) Ex.zi 1.tt (z', Z') exp (a llnes 1(=", x x1-1-h))
where 7 (2-M..... 2-1) and μºt the stationary distribution of the Markov chain over a window of size M + 1. We denote by pix (ex) = xx Ea) where Pixx (z) is the marginal distribution for
and serves the population counterpart for 1y Et a7 We define quantity A and B as
(EL+1 = =x)=
y (k) (01+1= ex)
A:
Σ
Σ
II hEST
1(x =11+1-
11 (k)+
(k) (k)+ +
M=M+1
E
1 [xL+1 = 21 =x)=
1 = exs)
II
1(z 1( ((0--==
(x
-=M+1
1=1
hES-
Under Assumption 3.5, we have
8(1 - A)-12 (No + 8/M 2de
1)114
14 - B <
L1/2 11111+11
70

---

Proof of Lemma F.10. The proof follows the same arguments as Lemma F.4. We remind the readers that y² (k) is also a function of the whole chain X. We note that
1 = =I 10== =xx)
[2 = =x= cx)
LA AA-B
B
E
Σ
oi
(*))+ +E
(x((EN)
-1
M+1+1
kE(d
kEId
y (k) 1 (IL+1 = ex) +1)
II
(xrs- = =
ya (k) +E
kE(d
hES
E
Σ of
E
(ex - y- (k)
(v)( + +6) (xlle))
E
1(11+1= = =es)
(y*(k) 8) Mx (ex)
M=M+1
kEld
E1(0141 = ex))
II 1(01n==ILHH
y*(k) +
kE
LES*
To handle this enror, we define three error terms
(Ex( - (k) +6 (ex)
Σ of (01+1==1-0 = = II
1(m=n= =(
erT ==
kE(d]
I=M+1
hES
Σ of 1(1111 = II
1a
err2 =
(1" (k) + 8) 1xx (Ex)
kEld
I=M+1
RES
L
E
Σ
of II
1 = r2+1-N)
CITy =
(k) +
kE|d]
Im M+1
BES.
For the first error term, we have that
(ex)
of 1(- = =x)
erri E
Σ
(v-(k) +))
(xx(
1=M+1
(ex y5 (k))
(kk)
1 (cw) - x" v(+1]]
E
E
(g)(k) +
Xx (ex)
[E[d]
kEId
where we recall that by assumption, 7 provides lower bound for x- Xpa), hence also lower bound for xi (ex.. Next, we invoke Proposition F.19 which provides an upper bound for the difference between the empirical and population distributions in terms of the -noor::
4((1 - A) DDa
110)) (DDlA)-VDeaalull
4M) 1/2
(==== = -
11/2. min, 1 10 (IL411
4(1 - 1) (D2 (=)) 1)/1/ +8V
(F223
L 11/2 1811+11
Hence, we control the first error term.
For the second error term, we follow the same procedure and obtain an upper bound as
Σ of (m===
err2 E
E
-de. de.
x (Ck)
(v*(k)+
P (xx)
kEld
I=M+1
kE
71

---

For the last error term, it holds that
crry
S
E
y (k) - 1 = =xa)
ke(d
E
1 - =
F (ex) + +
+
Σ REid]
E
(k( (k)
-xx(()))) - (7x +
( = =x
+) (cx)
+k)
kEld
C E
(v((k) - (ex))
4(1 A) (D2 (14 (²)) + 11) + M
+
1/2 7/1
7
?
kEId
where the last inequality follows directly from (F.23). In summary, the difference between f2 and fa is bounded by
1/2 fal S erri err2 + erra 8(1 X) 2--
(D2 -
(140 11 + 1) + 8VM 2de
1,1/2 1-11+1
which completes our proof of Lemma F.10.
The following lemmas are for analyzing the error Iifs fal for Stage III. Lemma F.1. We define
A A==Bi E
Σ
oi
1 (x1+11 = =1 =a) (x1+11 (1L+1=Ex)
((x)1==x))
=
II I (01 -
(x(()))
I=M+1
k=1
BEST
B: B:=Exx(x,2)xPP
1( = =2= Cx)
1(11+1 =ex) -ex)))==)))))))))))))))))))))
II 1 = =L(11-) LEST
(ex((x)
where
exp (a Ihes
of E exp (a I [1 = x1+1-i1)
(2, 2) exp (a lhes 1(2n = =1+11-h)
Xx (z, () pr' (z, Z X1+1-s-) Ex.21 1rth (z', Z') exp (a llnes 1(2'_, = x141-1))
Under Assumption 3.5, we have
A A-B|s 87 (1 (HO |1 |²)²a²) + 11/1/4 + +16---166-0000
11/2 715-1+11
Proof of Lemma F.11. For 2 = (Z-M,..., 2-11 and Z' = 2(1)) we let Z=s-
(z-n)kEs, we define
1
(x 7) -- M
Σ
11t XI-MM-1 2),
L=M+1
R(2,X exp
a II
= ((=-=0ot1-1)
hESt
Using these notations, we can rewrite the normalizing factor in 7 and of respectively as
+ ΣnT (2, R(Z,X1+1-5)) き= Ex 2) ((
2
2.7
We also define
662, 2-s+) = Z-s-): (.Z=s.,-s-)).
= (.
72

---

If we further define (2(2,Z-s:) = =Ei=M41 (22) X1 Z-s+), then we have
(2, s.) 1x (2, -.)
6(2, Z-s.)
(2, 11" (2,7- s.)
(2,
+
Using the above definitions and relationship, A and B can be rewritten
d Hee XL+1-S-) =
2 ((X++1-s))
plek, X
x(X
B E
(x))(e))
$
Therefore, the difference between A and B is given by
2 E AA B S 7
Σ
(2,,7_5))
d(z, 2.)
)+)
__.B
2
Σ
-
s' S-
8y-1, ((1 - 1)- D22 (10 (a)) + 4M)
1/2
1/2 min SCtT
where the last inequality follows from the result in Proposition F.19. Invoking the lower bound 1² (IL+1. 2 187|+1, we complete the proof of Lemma F.11.
F.3 Lemmas on Concentration of Markov Chain
Recall that we previously define X (x ..... TT) as the observed sequence and x1+1 the value at time + 1 to be predicted. For generality, we will use X = (Z],..., F1+1) to denote the whole sequence in the following proof. We denote by p" (-) the joint distribution for the sequence X with
kernel и. Recall that we have the parent set pa = { (-r,...., -ra), and as the start of a chain, we
sample the first tokens by (ZI,..., (r,)) M40
In the sequel, will study concentration properties of the Markov chain X for a window of tokens with window size at most M, where M > Tn- To proceed, let us consider a fixed set C [M]. For any 1 E [M + 1, L + 1]. we define Yi = (æ1 Xi-s) as new vector containing the token at position and also the tokens in the past S positions prior to Hı. Here, we follow the convention that Xi-s (Xi-)kes. We also consider another fixed subset S' C [M] and similarly define Y/ (x1,
The concentration properties of the Markov chain are rooted the fact that when conditioning on all the parents, the current token is independent of all the past tokens. Given the parent set structure pa = {-r1...., } we aim to make YL+1 approximately independent of Y; by conditioning on some intermediate parent sets. To this end, we define A = (IL+11-M
#LL-M++ +r.)) ) E 2 and
B, (II-rut1..... x1) E XT these intermediate parent sets. By the Markov property and the parent set structure, we have the following conditional independence relations:
YL+1 TL (B,, Y)) AY)+1, TL | VI M ... L M + Tn-
To illustrate, let us consider the first condition YL+1 IL (B,, Yi) A. When L M + r.. the Bi and Yi are both contained in the history { k < L M nn AU {xx: k { < L--LM M). When conditioning on A, the randomness of (B,, Y)) measurable by the oaaleerra generated by the "past" {zx: & < L M). Moreover, the randomness of YL+1 is measurable by the o-algebra generated by the "future" {aRR: : l E [L+1- M r,, L + 1]} when conditioning on A. Notice that the parent to the any element in the future {IA: k E [L+11 11 M + +T., + 1]} is either contained in A, or can be generated conditioned on A without touching further history {Ix: k $ L M). Thus, by the Markov property, conditioning on A, Yi++1 independent of the past {ux: k: S L M). and in particular, (Bi, Yi). Similarly, since B contains the parent of 2/+ conditioning on B, Y, is independent of æ141 and later tokens. Moreover, given B, the randomness of Y. comes from the randomness of x ......... TI-r.,. Since < I M + Tn. we have I. + 1 - M 2 11+1 11 Ta- As result, conditioning on B, the randomness of (Yr+1, A) comes from tokens generated
earlier than Z1+1- Therefore, (Y A) and Yi are conditionally independent given B, We visualize the definition of YL+1. A, B,, and Y in Figure 10
73

---

2 1
Y,
X++1 Mtr,, Xt-M 1
X X X/-r,
A/-M
A
B1
Figure 10: Illustration of the definition of L A, Bt, and Y,. When conditioned on A, Y1-1 is independent of (Bt, Yi). When conditioned on Br, Yi is independent of (A, Y1+1).
Similarly, for Y/ (Ii, si defined using the subset S', we also parallel conditional independence relations:
Yi IL (B), A, (Yi+1, A) Y Bi, VI M + ,....,...
,,, L M + Tn-
In particular, we also have
YL TL (R,Y)) (Y4+1, 11 | = M+ L M + Tu
(F.24)
Using (Y, Y.), we define a joint distribution 5 over 2 + |S| + SS' tokens as follows. For any E' E x1S++1 and E' E x114+11 the probability mass function of pr is defined as
(VL++1 = E, = F')
1
1
Σ
EE)
L M
1=M+1
L
EX = A) P pl- M-T, 11r--1(60 (A| p (X' = EE) . p" (B:).
L - M L A.B.
(F225)
Here, y" is just placeholder for Y takes average over and does not depend on any specific position index. The summation ΣA,8, means we sum over all possible values that A and B, can take. In the last line of (F.25), we decompose the joint distribution (" E, Y = E") into the product of the conditional distributions by the Markov property in (F224. That is,
pd (YL-1 E,Y E') Σ
(Y E,Y, = Bi)
A,
Σr'
(Bi) p0 1" = E, Bi) PV EB)
A.B:
= Σ (Y E ( (A Bd) p (Yi = E' Bi) (Bi).
A.B:
Here the second equality follows from the fact that (YL+1, A) 11 Y/ B, and the last equality follows from the fact that Yr++ JL. (Br, Y/) A, which implies p² (Yr.+1 = E A, Bi) = p" (Yr.+1 EE A). Moreover, we denote by Pt the i-step transition kernel of the chain, which corresponds to the i-th power of the transition matrix P.. Here, we are following the convention in the main text that
(Z Z) π(2 Zpa(t) - 1(Z Zi-r. +1--1.
(F.26)
In the following, we always consider fixed transition kernel and omit the superscrip//subscrip in the matrix notation. We denote the transition matrix by Px and the stationary distribution by jr"
for window of length r. For the transition matrix, we index each row by the next r.- wwndoww Z' and each column by the current r, window Z. Under this notation, since both A and B; have lengths we have
p" (A Bi) Mtr - (A, Bi).
(F.27)
Here pf. pL- Mr, - ((A) B;) corresponds to the (A, Bi)-entry of the matrix (P₂) M -1 Combin- ing (F.24) and (F.27,, obtain the last equality in (F.25).
74

---

In the sequel, to simplify the notation, we write P7 and μ= as P and μ respectively. Let us consider the reweighted transition kernel
k= diag (v V=)
P. diag (VP),
where v is the clement-wise square root of μ. Since the transition matrix is primitive by assumption and having only one eigenvalue with value one on its spectral circle, we also have for K that the leading eigenvalue is one with eigenvector vfl, i.e. vIi K Ky and vi / K. However, the projection in the leading eigenspace (or the Perron projection) is not our interest. The following property of K will be useful in the subsequent proof.
Proposition F.12. For the reweighted transition matrix K, we have for any integer 2
pi pl T diag (VT) (K vTVTT)'* diiaggv (vT=
Proof of Proposition F.12.
pi-p1' (diag( VTl) K .diagg( (VTg((Vi) " 111
diag( (v)) (K' (K'- VHVV diag (V
diag (Vπ) (K- VVVV= diag(
(V
where the last equality holds by noting that K - VuVT" project v to the zero vector, and for any 1 Il, have (K /PV/ )u= Kv. Thus for any test vector
(K - VRVK y'x =K (K-
(VH,z)) VT) = =K'(x (VH,z) vp) =(K" VV-VVV
VaVV y(x -K-VVuV
This completes the proof of Proposition F.12.
Indeed, the second largest eigenvalue of K (in magnitude) determines the mixing rate of the chain. Let denote the eigenvalue of K with the second largest magnitude.
Furthermore, if the transition kemell admits a lower bound y 0, then can guarantee that both p" and μt admit uniform lower bound.
Proposition F.13 (Uniform Lower Bound). Suppose T(- Xpa) > zy uniformly for some 7 >> and pa = {-r1,.... >ra}. 5 Suppose X1.r. ~ (o(()) where μ4 E ((X"=). Then for any S tokens
Ii ......., Fig such that lg Tn for any E [S), we have
p" (Z)g...., m1g) 7.
Using Proposition F.13, we show that the transition matrix Pa is primitive.
Corollary E.14 (Uniform Lower Bound Implies Primitive Transition). Under the condition of Proposition F.13, with T- Xpa) 2 O, the transition matrix defined in (F226) is primitive.
Proof of Corollary F.14. If the initial distribution is set to be any one-hot vector in A(AT)), and taking p= ("rr. X 1irr)) 0 holds for any :
Th...... Big in Proposition F.13 to be fr..1, T2r.. we conclude that
E X"a, Recall from the definition that for a primitive matrix P, we can find some positive integer & such that pk has all positive entries. For our case, we can set rn and everything follows by noting that pd (X,,+1:2 (X1rr.) pro (Xr,+1:2r, 1.r.))
Another corollary of Proposition F.13 is that, if take μaa 1x",, which is the stationary distribution, we can replace p" in Proposition F.13 by μtt
Corollary F.15. Suppose T(- (-1xpaa) 2 y uniformly for some y > 0 and pa {-ri...., -Tn) For
the stationary distribution μr² and S tokens Fiy Tig...., Fig such that I, 2 Tn for any E [S]. we have μº (II, > S.
We prove Proposition F.13 1133a. as follows.
75

---

Proof of Proposition F.13. Without loss of generality, suppose that M < 1 < l22 <... ls. We will prove the statement by induction on the number of tokens S. If S we can rewrite
r" (m),
Σ (,,,, Xxa(ti))) (Xp((((())))
2
Σ r.p" (Xpp(t))) 2 7-
Xpsl())
Now, suppose the statement holds for 1,2....,S 1.1.ee - -t,- Then, have
p² (II ris) =
Σ
{(āis| (xa(ls)) p²))))) p" (Xpalls) \
X
Σ (VP p² (Xpatte) \Y) = 7 7. p" (Y) 2 7
XpussY
where the last inequality holds by the induction condition. Hence, we finish the proof.
Before analyzing p", we first study a simpler convergence result: quantifying the closeness between Ei-M+1 n/ (B, = b)/ Ei=M+1n' and μ (b) for certain values of n E (0, 1].
Lemma F.16. Following the notations introduced above, for the Markov chain with parent set let Dx2(Po I (²") be the x² Xi-divergence between the initial distribution po and the stationary distribution μºº over the first I'n tokens. Take any S C [M] and let Yi (zi, Xi-s)
pa = {-ri,-..,--y
for = M 4f+1.. 1,..., L 1. Suppose I.22 2 M 2 Tn. We have
E (Vi
2vDz² (po 1 1 +1
10 (YL++1= -1)=
I
(F.28)
L-M
TV
L(1
(Y - 1 (Y(+1= = )|Tv S X--M M Dza |1 (17) +1.
(F.29)
Proof of Lemma F.16. Let Let E
T- tr. n where 17 E [0, 1] is constant to be deter- Lizr
mined. Denote by po. a vector of length 1x""", the initial distribution of the chain. We begin by quantifying the total variation (TV) distance:
Eis ELaM+t X (-'D(((==))
ΣE LTーM-T X - 1²(() --
L- M-M+r
= Σ a. (B)
TV
TV
Let b E X",, representing the value for length-r, window. Using matrix notation, we have:
M+M
L-M+ro
L-M+ra
Σ 1=P
a (B) = b) (b)) =
Σ
a.1, pl ((p - -1)= Σ a 1 (p'-r - -m1? nm
I=rn
1=
L- M-M+ra
Σ
a. 1. (VP) (K () diag (Vr) p4.
where 1, is the indicator vector corresponding to b. The last equality follows from Proposition F.12. For any test vector u E {0, 11111" using the variational representation of TV distance:
M4
+rn
Σ a (B)
max 1Ef0,1)ix1na
TT
Σ
q (B:
TV
L-M+Ta
Σ a. diag (VR) (K (MVWT)"
diag (v -40
max
L- M4
L- L-M
<
Σ
q 2.
diag (VHT) -10 -10|1 = Σ
a. --xi / D22 (1F) + 1,
1=
I=
(F.30)
76

---

where the inequality holds by ||u diag (VH) S IVAl2 = and K-
- VAVV has leading eigen-
value with magnitude A. The last identity follows directly from the definition of the x² diverrggen that Dx² (No ll (²)) + 1 = Es Po(b)) ((0)))//((00000 (b).
Substituting the definition of ct, we have
+r (B
L-M Ii birr
mptl X-TT
DDx (1 +1
1 (A- b)
Ei=ra El n--
Ei=ra El-M+rn np
TV
We consider two special cases. In the first case, we set = X, which gives us
El-nttr (B) (B)=b)
L-M+ra Nb- DDz 1C) +
Er-Mtt A6-1 1 1= =1 b)
Limro
(1 (1-L--M) (11 -))
TV
L. A ro (1-
VDx (1 ( +1.
+1.
1 X1-M
In the second case, we set 17 = 1, which gives us
M +Tn p" (B) =) 1 (A=-)
DL-M+Ta X-ro v D22 1+++1
D2² (10 11*) +1 (I. M)(1 A)
L-M
M
TV
Note that the TV distance is an f-divergence. Thus, we can use the data processing inequality obtain
the desired result for Yi from the above inequality. To do so, note that _/= P" (Vi = -)/ M)
and n= (YI.+1 :) can be transformed from Ei=rm the same emission kernel
Et -1 p" (Bt )/ M) and n² (A= by
p" (Vi+1 = . A = .) pV =: M-r, = 1 (VL+++== = | )) 10 (V, - B B--M+T.- =))
Therefore, by the data processing inequality, it holds that
-- M 10-10"(M1 :)
1= 1=t(""n (VL+++= (V+++++)) +>)1=>>)
Lizra trn (B1 =.)
< VDza(po 125))
1z (A(A=)) (A=.)
M-M
TT
L-M M
TV
(L - M) A)
Similarly for (Y(.t1 :) and (-), we have
(Y(+1=1)) - p² (Yr+1 = l|Tv S I1p = 10 1²"(AA = )|TTV
max TT -ddag (VK) - VVVV ) L- diag (v²) -1 Po S X NE(0.
M
M V (mo 11 (²) + 1,
where the latter two inequality follows from the same arguments as in (F.30). Hence, the proof is completed.
We have established that the average Σi-M+1 -Y; - M) converges to (A = .) in total variation distance. This represents "first-order" convergence since it involves the average of the marginal distribution of Y. However, the quantity of interest in (F.25) is the average of the joint distribution of Yr.+1 and Y, which concerns "second-order" convergence. This is studied in the following lemma.
Lemma F.17. Following the notations introduced above, for the Markov chain with parent set pa = {-r1..... -rn), let a (Ho 1² be the x² -divergene between the initial distribution pn and the stationary distribution μ over the first In tokens. Take any S, S' C [M] and let Yi = (z), Xi-s)
and Yi = (Ii, Xi-s:) for M + 1,.... L L+1. 1. Suppose L/2 2 M 2 Tw- For f defined in (F:25),
we have
2M
4VD Dxx 11 (1) +1
T (YL+1 pº (YL+1 = (X' -llvv &
+
L
L(1 A) mine n (Yi+1 E)
In particular, we have
(YL+I (YL+1 -10 -
Σ
pT(Y( =
1 M
TV
1= 1=M++
2M
2 DvDg2(11 +I
<
+
(F31)
L
L(1 - A) ").Viinng /ming// 1º5 (YL+1 = E)
77

---

Proof of Lemma F.17. Let us take (E)E). (L - M) -EEE--I"PP E") as the intermediate distribution, and we have by (F.25) that
Σl.
1
(Y = E,Y - (Y++1=
Σ
p pYi = E3))
M
i=M+1
1 L-M+ra
L-M Σ
Σ
1 (VL+1 B1A). (pl-I (m) (A) Bi) 1= DV' E' EB)B)) (B.)
1=M+1 A.B.
(I)
1
Σ
(Y E,Y' = - = - ='))
(F.32)
L M
(II)
where we use the fact that Σ> term on the right hand side can be easily controlled as we already have an L-1 factor. We let TV, be the total variation distance of the second term. It is easy to that
A 1e (YL+1 E' (A) IT
E) for the first line. The second
M - In L M
M
TV = 2
Σ E,E"
ICD)
S
L-M
where we remark that (II) a function of both E and E', and the total variation distance just taking the sum of the absolute values of the differences. Here, we also use the fact that L 2 2MM Using Proposition F. F.12, we can also rewrite the first term on the right hand side of (F.32) in the matrix form as
Mrn Σ
(YL+1 - diag (VØ) (K )
(I)
diag (vū)
L M
M=++1
diag(p"((B .p (Y = B: T.
When considering the normm of the above term, we introduce test matrix U of shape | x | with each element of U chosen from {0, 1]. Let TV, be the total variation distance of the first term (I). Then, we have
L- M+
TV, & max Tr
Σ
(VL+1= .diag (VT) (K])(KK-))) - VTVTT r. 1. (M5-))) (MM-)))
U
L M
1=M+1
diag (Vr) diag(p" (Bag(p"(B:=))) ----)))>p'000 DYi = - -) -
U(..)T.
To upper bound this quantity, we consider each row of U as U(E,)) = u(- T) Note that u(- E)
is also a (0, 1a-valued vector. By expanding the trace, have
L- M+ro
TV S
Σ max M(-1 E)
Σ 1 (YL+1= =EA==) diag( (VT)
L M
1=M+1
(K VHVR (- diag (VT)-1 diiggpp (B) = -)) -p' = B :) -u(-1 E).
Note that the f2-morm of the vector in the last line can be upper bounded by
(K- VMVH 1--- (M diag (VH) diag(p" (B) = .)) pd (Y' = . Bj T). u((- E))
||----- diag (vØ) -1 diag(p" = .)) 1 = X 1- -1 (M- ||daa (vZi) (B:(Bi = =
(M
-M--)-(M-r)) / /D2 = -1 ))) + 1 (- / DDD (14 (15) +1,
(F.33)
78

---

where the first inequality holds by noting that p"(Y?' = I =) -u(-EE
is vector with
element within [0, 1]. and also invoking the operator norm of the matrix K v The second
identity follows from the definition of the x²-divergence that Dxa(x" (y=(""B, (Bi -) n (+) + En (Bi =)] //=" (b). The last inequality is the data processing inequality p" (B = can be transformed from (a Br..) and μ=" (Bi) can be transformed from fz (Br,, by the same emission
kernel μº' (Bi = Br. .). Consequently, we have for the TV distance that
-- M+
T
1
Σ
X (M-T-( V Dx (uo (T +
L M
M+1
max
EIAA) (A) - ((A]E)
fo(. | E)}e:
[ B E,A
DDx 1+) +1
1 -++=-E)
max (u(- B)) ll 1
o(A E) pp" (YL+11 =E]
M)
(A*(A)
A,E
VD D2a 11) (T. - M)- (AX+)i1) +1.
< [Mt-) (E)E (S)))
[A] (Y E).
A.E
where in the first equality, we the variational form of the f--norm for vector 2rt (Y 1+1 E .) diag( VP). In the second inequality, we apply (F.33) and use the Bayes rule. The last inequality follows from the Cauchy-Schh Schwwrz inequality. Here, the mutual information (A;; (L11) + can be
upper bounded by
1 (Y = EIA), (Y F,A) ming μ (YL+1 E) A, 1 (Y E)
Tx (A):
and the last term involving u(Al E) can be upper bounded by 1 thanks to the constraint on vi. E.. In conclusion,
Dx (no (x) + TV S (L-M)(1 A) / /ming feth (YL+1 E)
Lastly, let us relate the intermediate distribution to the final distribution is (Y =) x 1r x) where we define the total variation distance TV2 as
=)).
TV2:
TV; - 110 (
(
Σ 1 1 on -)
=
(
1 M
Σ
-( -17()
L=M+1
TV
i=M+1
Invoking (F:28) of Lemma F.16, we have this quantity upper bounded by
VDx 1+)
TV S (L- M)
Using the triangular inequality for the total variation distance, we have
1 (YL+1 = =,, :)- n² = =
-VV+TVV TV2
M
2 Dy2(H0 1TT) +1
+
b M
(L- M) 2) /inu nrt (Y = E)
2M
4/ 4VDz2(H1 10)) +I
L
L(1- /ming 4 (YL+1=BI
and the upper bound for (F.31) follows by the same arguments. Hence, the proof is completed.
In the following, we use a similar technique as in Lemma F.17 to derive a bound for the chi-square divergence.
79

---

Lemma E.18. For the x2 diverrgence between the empirical distribution - -M-- E|=M+1
1(Y = and the stationary distribution je (-), we have
.nd
E Dx² (T. (
1 Σ (M)=))|/~(Yann
v 4(1- 11- /D2a20-VD 11))) + 1 1+ 16M
L ming μttt (YL++ = E)
1=M+1
where the expectation is with respect to X ~ p=.
Proof of Lemma F.18. By definition of the x2 d'dveergeece, what we aim to bound is just
L
E
(L M) -1 Σ
10y E) 10 ( 10 (B)
E
1= M+1
(L. - My
M+1==++ 1(%- = Yr = E) - 1B9
E
(F)
=E
Σ Σ
Σ
L
1( Yv E)
1 = Σ
Σ
pT Yp E)
L11=M+1 (L- (E)
(L- - (E)
E 1+1=M+1
To study the above quantity, for 2 2M - -nn + 2, we define
J1(1)= Σ Σ
M-M+r.-1
(Yi Yv E)
1- 2M + In
(L - M)²
M4=
Following our convention, we let A XI-Ma- -.-MMr.------ -1 and By X1-ra+1:1' be two length-r, window and by the Markov property, we have
Y1+1 1L BBrYYrr A,, BI) TT Yv Br.
Let us fix an index 2 2M In + and take a summation over M + 1 S 1' 21 M + rn 1. Expanding the joint distribution, we have
1-M+r.-1
J²(1) = (L - M
Σ
Σ
= (VN=E|A) (p'- -1 (Ail B:) - -(r))(((A))) (Ai)
M=
B,A.,Br
p" "Y,
=El E|
. 10" (Yp = [{)=
1- L
Σ
Tr (ii((i = AAi -diag (Vr) (K - vRVzI -
(T. M)²
M=M++
diag diag(vR) adiag(p"((B .)) (Yv = By =,) diag(u" (Yr- = -)-1
1-M+
Σ
Tr diag(u (Yr = 1/11/) 1¹ 10i(ii . A1 .) ..diag (V11)
(L- - My²
#1
(K VPVA T) - diag( VH) diag(p" (p" (Br ..)
-Y= = =V By -) diag(n (1t (Yr = .) -111/2)
where the first identity follows from the fact that
Σ
H¹² * -Y Ar) 4º (AI) p" (Yr- = E Bl)p .p (Brr) 11 (Yr = E)-
E.
= =E DY' E) μt (Yr- = E) "Y = E1-1=1 1,
and the second identity follows from Proposition F.12. We next invoke the Cauchy-Schware inequaaly for trace, i.e., Tr(W V-² < Tr(W W)ThWW"V Tr(VTT(T]) where we take
W = diag(H ( (Yi = -1/2) (Yi = -1 AA = -) ) .diag( -K VTVTTT ) -M V diag( (VTi) -1 diag(p" (Br- 1) ---((Yr (Yr = 1 =)) diag(n" (Vr = -1-1/2
diag(v/) "Y² = Bp = T) -diag(u ((V -11/
80

---

Note that
(T<(WTT) W) s (----M
(diag((* (n* (Vi = - (Y== - A .diag (1) 1 (0i - -AA
(Y),
-
A,,Yi (Vi) 10 (Ai)
Following the same calculation, we have
Dd (Yv, B 10tt (Yr))a (1 (Be)
(T(VTV)=
Σ
By,
Therefore,
1-M+ra-1
Σ HY,,,, Ai) Av.Yi (V) (Ai)
PYr, Bv Σ
Ji(1) S (L- - M)2
Σ
X -1
YerBr. 10t (Yr)) (Bu)
1'=++1
We further have
(Y, B²
pT |(Y) Br)
(Br)²²
Σ By
max By
(Yr )100 (By)
11. (Yv)
1ttT (By)
By
Pra(0" (By =.) p1 )) <
Dz2 |l 11*))
ming μ² (Y = E)
ming nM (YL+1 E)
where the last inequality holds by the data processing inequality. Similarly, have
10 (1) max 1 (Vi) (A1)
[Ci [[))
Yi,A
1=t (Yi)
ming μat (Vi+1 = F)
Therefore, we conclude that
/ DD (no 12x)
(1 E
(L M) A) ming nmt (Y E)¹
and
2 D2>02pppp (P) (T)) +
2
Σ
J1(1) S
(L - M) A) ming nm (YL+1 E)
2M
where we double the value as 1 >> 1' only contributes to half of the terms in the double summation. Note that the above summation for 1', we only include terms satisfying I- M - rn + and - (M + 1) 2 M - rn + 1. For the remaining (1, (') not included above, each term is bounded
above by
pd pY (Yi
Yo = E) 1)
2)
10 (E)
(L M)2 ming pm (YL+1 = E)'
E
and we have more than 4L(M In 1) of these terms in total. As result, we conclude with L/2 2 M 2 rn that
J 5 4(1 - DD= 10)) +1 + 16M
L minz 4tt (Yt+1 E)
Hence, we complete the proof of Lemma F.18.
Proposition F.19. Let define
1º Z- exp (a IInss- 1(2-1 x2+1-11) A3 2.. = Ez 1r (z', Z'-s.) exp (a Ihess 1(x"_, = xI+1-1))
81

---

where S- (2-N)lEs- and μ²' is the stationary distribution of the Markov chain over a window size M + 1. We also treat (-)()) as a length |X] vector where X is the state space of the Markov chain. Let 8% (z, Z=s+) = Ei=M+1 of 1a: = 2, XI =
=Z-s+) where
exp(a Ihess 1 =
of
Σ M+= exp(a IhES (xr-- #1+1-1))
Then, we have
- 2. = - 1,1/2 ( -(11-))-VVDxa (()) +1++1111 + 4M) 1/2
=
-)ll
-1 1et (IL+1, X
Proof of Proposition F.19. To unify the notations, we let Z = (z-M, ...,,.... ,-,, and define
(x =
Σ
(( = 2, X1-5- -._s)),
M=M+1
R(Z s XL+1- = exp (a 1(2-5. = 11+1---)).
Using these notations, we can define the normalizing factor in Pa and 1x respectively as
Σ (2, (ta-ss)))))) $ = Σ
xx 2-))
:
2,2 s
We also define
652 2-s+) = (2, ((2, 1
= ( 212
We can then rewrite the objective as
x (= 7- =
(= =
0(2, z.)
Z(2,
Z(2,Z_5.).1 1 + 2 -0(z Z s1
Σ
+
Φ
x.7_8*
2,7_80
@ 1 Ex,7.g gr b(z,Z
((2;Z_5))
2 10(z, --) ((2, Z-s+)
+
Furthermore, notice that
Σx,2 sa (6(2ZZ-s)) Z-2))
Σz.z ar (2, - 1
Ee,h 11 (2,b-s-)
Σxt (H, 2-s-) - -2x (((2,z-=s)))1 + (c" - 1) 11 E,,),,,,aa
S1ET 2
22
(2, Z-.-) - (2, Z-s+)
1 (e" -1) E,.z..-ETxx (2,
Σ (z, - ΣxIZ.seE's 14 (2,7---) - - -xx(2, 7-s)))) 2,7_8*
ExZLo Hº1 (2, Z.s
(F:44)
where we define Tx = {7- 7- = -14+1-s+} Note that when 7-5: E [xx we have R(Z R(Z-s", XL+I-S-) = e and when Z-s- Tx. we have R( R(Z-s., Xi+1-S.) 1. For the first term on the right-hand side of (F.34), we have by Cauchy-Schwar that
(1r (z, Z_s - - (2, 7 s1) ))
1/2
Ex
Σ
Z z
<
Ex
n (2,
2.7_8*
2
4(1 X) D 11
1 +1+ + 16M
1/2
(21+1.
st
82

---

where in the last inequality, we invoke Lemma F.18 where we take Y = in the lemma. For the second term on the right hand of (F.34), we note that
Ex.Z_o.
ETx 1(0" (2, (-) Pi (z, - 1>l
Ex
Σxit stElx 11 (-)
11" (2, Z-s. = =E) - 11x (z, Z- = =E]]
1(X = =B)
Ex
(7 = F))
E.5
11 (2, Zs s. = E) - -xx U- =b))²
p² = E3) 1/2
Σx
v V¹T (7 = E)
11t (7-55 F.)
(1 (2,66_. = E) - 11x (2, -- = -)))
p" (' = =11) 1/2
1
Ex
1r (Z = E)
p² = E)
F,2
E,E
(F.35)
where the last two inequalities follow from the Cauchy-Schward inequality. We have an upper bound for the second term the right- higtt side of (F.35) that
1= = - = E)
p"
1/2
1
ming 1t (Z-s- E)
We can also apply Lemma F.18 to the first term with Y 1 X and conclude that
(n* (z, s = E) (2, 2- (2, = E))² B))²)
1)
1/2
Ex
1= 1= E))
E,=
4(1 V || 111 + 16M 1/2
-iilz 11t (1141)
In summary, we have
Ex (ex) v/' (k)
(1 (1-A)-1 v DD2 ) +I+ 4M 1/2
2
minert 1ºt (I1+1
(1 ( 1 +1+M
1/2
L --miingt 1*t (1L+1:
Note that the second term dominated by the first term. Thus, we conclude the proof of Proposi- tion 1.
83

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The convergence results are established in Theorem 3.6 and we relate the learned transformer model to the generalized induction head in the following discussions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed in $B.2
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
84

---

Justification: We provide the full set of assumptions in Assumption 3.3 and Assumption 3.5. We provide a complete and correct proof in $E and $F.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Ys]]
Justification:
We provide the complete details for our numerical experiment in $B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the reprrcuuuuee is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
85

---

reproducible, they Answer: [No]
Justification: We do not release the data and code, but the details provided in $B are sufficient for reproducing the synthetic data and the experiment results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips cc/
oolliccyy for more details.
While we encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips.co httpsccccpbbcccccccccccccccccccc wttppppppwiddss.//p/pppppppppapp SStm//ubdeSiiccacaceeeee ississsiii sassnoolicc) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide all the training and test details $B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We do report error bars because the training behavior is consistent across different runs, and the goal of the experiments to corroborate our main theoretical results, rather than achieving better performance benchmarks.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
86

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We provide information about compute resources in $B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: hww./erursp...// tttp:::/aurrpp..coobllttt/////// Caaublcc/Ehicsuuute
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and we confirm that our submission adheres to the guidelines therein.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make ation due to laws or regulations in their jurisdiction).
preserve anonymity (e.g., if there is special consider-
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: In the current paper, we focus on developing theoretical understanding of transformers, and the goal is to analyze existing architectures instead of proposing new models for better performance. Therefore, we not see immediate societal impact of our paper.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
87

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not release data models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: We use common and standard Python libraries and write our own code for the experiments. Also, the data used in experiments is synthetic.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswi paperswithcode thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
88

---

13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not introduce new assets in the current paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: We perform theoretical analysis and numerical simulations on synthetic data, and the process is not related to human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: As clarified in the answer above, our study is not related to human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
89


Reviews:
[{"review_id": "3mY3fyAPPR", "reviewer": "Reviewer_omVG", "summary": "This paper theoretically explores how simplified transformers perform in-context learning (ICL) on n-gram Markov chain data. The authors analyze a two-attention-layer transformer model with relative positional embedding (RPE), multi-head softmax attention (but second layer only scalar trainable), and an FFN with normalization. They prove that gradient flow concerning a cross-entropy ICL loss converges to a limiting model that implements a generalized \"induction head\" mechanism. This mechanism involves copying past tokens within a window and selecting informationally parent tokens using FFN, followed by classification through the second attention layer. The theoretical findings are supported by numerical experiments.", "questions": "Similar to Weakness. There are some other questions:\n1. Using the second attention layer as a classifier is a little weird because people always use the last linear layer to play this role, but in your framework, you delete the last linear layer. Considering the parameter of the second attention layer is quite simplified, I doubt that the second 'attention layer' is necessary (like we don't need the attention token similarity) for getting your results. Is it possible to remove the second attention scalar but use something like a linear layer to classify?\n2. Is line 247 $\\psi_{S^*}(i)$ a typo? Besides what's the definition of $i_s$ in D.3? and what's the dimension ?(also a typo here?)", "limitations": "The authors don't analyze their limitations", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow\n2. The definition/assumption is clear. The theoretical framework using chi-square mutual information is solid and novel, within which the authors can determine the most suitable parent sets that are both informative and not so complex for handling the n-gram dependency. The conclusion of the three-stage training is also interesting. The process is theoretically provable and intuitionally clear, like first training FFN to find the most suitable parent sets, then training the first attention (RPE) to select each parent for each head, and finally tunning the second attention scalar ($a$) to serve as an exponential kernel classifier.\n3. Use empirical experiments to support the theory's results.", "weaknesses": "1. The setting for the attention layer may be too simplified. For example, the authors don't consider residual connection, and they also don't consider the token similarity (he original W_K, W_Q), but only the relative position embeddings, i.e. RPE, in the first attention layer and trainable scalar $a$ in the second attention layer. The training process is manually split into three stages (first attention RPE-> FFN -> second attention scalar $a$). These may make the problem more theoretical-friendly but quite different from the real-world setting. \n2. Following Weakness-1, the numerical experiments are more on synthetic datasets + simplified models rather than real-world datasets (like wiki-text) + general 2-layer transformers. Thus I doubt the potential of using this framework for more general transformer architectures, although this problem is indeed quite difficult.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. Please see our response below.\n\n**On the model simplification**: \n- *(Residual link)*\nWe want to clarify that the residual link is indeed contained in our simplified model via the use of disentangled transformer. As can beseen in Eqn (2.5), the second layer’s attention uses the original input $X$, which is copied from the first layer via a residual link. \nWe apologize for the confusion and will make this point more clear in the revision.\n- *(Model Simplification)*\nSimplifications in the model are commonly employed by previous work. For example, Bietti et al. (2023) also excludes the token similarity in the first layer of their simplified model, and Nichani et al. (2024) assumes symmetry conditions on the data which makes the second layer's QK embedding matrix as $W_{QK} = a I + b \\bf{1} \\bf{1}^\\top$, where the all one part $ b \\bf{1} \\bf{1}^\\top$ assigns the same score among all tokens and the model is essentially the same as only using a single scalar $a$. \n\nPlease see also our response in the general rebuttal on the data/model simplification and the split of training phases. \nWe remark that our model still contains the core components of the original transformer, including multi-head attention with relative positional embeddings, FFN, layer normalization, and also the residual link, which distinguishes our setting from the previous ones. \nEven with split of training stages, the analysis is highly non-trivial due to the cross-head interference and the nonlinearity in the model. \n\n**On the experiments**: Our work is theory-oriented, thus we mainly conduct experiments on the synthetic data and the two-layer transformer model to validate our theoretical findings.\nMoreover, as we have pointed out in the general rebuttal, the $n$-gram Markov structure is already more expressive than the bi-gram data studied in previous works (Bietti et al. (2023)).\n\n**Questions on attention classifier**: If our understanding is correct, the reviewer is asking if we can replace the second attention layer by a linear layer for classification. \nThe answer is negative. \nThe second attention layer plays the role of a classifier that first _compares the query's feature (which contains its parents' information) with previous tokens' features_, and then _aggregates the tokens with the same parents_.\nA simple linear layer cannot do this task as it requires the ability of processing sequential information. See Olsson et al. (2022) for more details on the induction-head mechanism.\nPlease let us know if our understanding is correct.\n\n**On the typo**: Thank you for pointing out the typo. Yes, $\\psi_{\\mathcal{S}^\\star}(i)$ in line 247 should be $\\psi_{\\mathcal{S}^\\star}(l)$ instead. We will correct this in the revision.\n\n**On the definition of the feature $\\psi_{\\mathcal{S}}(l)$**: \nThank you for pointing out the display issue in Eqn (D.3).\nIn the definition of $\\psi_{\\mathcal{S}}(l)$ in Eqn (D.3), which we copy in the following line, \n\n$$\n\\psi _{\\mathcal{S}}(l) = \\left(\\prod _{s\\in\\mathcal{S}^\\star}(x _{l-s}) _{i_s}: \\\\{i _{s}\\\\} _{s\\in \\mathcal{S}} \\subseteq [d]\\right), \\quad \\forall l\\in[L+1],\n$$ \n\neach $i_s$ refers to an index of elements in vector $x_{l-s}$ for $s\\in\\mathcal{S}$. \nHere, we have $|\\mathcal{S}|$ vectors in $\\mathcal{S}$ and each $i_s$ has $d$ choices. The feature $\\psi_{\\mathcal{S}}(l)$ is thus of dimension $d^{|\\mathcal{S}|}$ as it is defined element-wise by iterating over all the indices $i_s$ for all $s\\in\\mathcal{S}$. \nWe will correct this in the revision.\n\n\n**References**:\n- Bietti, Alberto, et al. Birth of a transformer: A memory viewpoint. 2023.\n- Nichani, Eshaan, et al. How transformers learn causal structure with gradient descent. 2024.\n- Olsson, Catherine, et al. In-context learning and induction heads. 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "38SgkAgANy", "reviewer": "Reviewer_NTvL", "summary": "This research focuses on analyzing in-context learning (ICL) in transformer models specifically trained on n-gram Markov chain data. It analyzes how the components of the transformer, attention layers and feed-forward networks, contribute to processing and learning from this type of data, demonstrating their roles through simulation experiments.", "questions": "Q1: Could you clarify the differences between this paper and prior work in the field? Specifically, the paper seems to focus on bi-gram models; how challenging would it be to extend this analysis to n-gram models? What are the main hurdles in scaling the approach from bi-grams to n-grams?\nBirth of a Transformer: A Memory Viewpoint\nAlberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, Leon Bottou\nNeurIPS2023\n\nQ2: What are the advantages of using the Modified Chi-squared Mutual Information (MI) for your analysis? Are there specific properties or new insights that could only be demonstrated using this metric? What unique findings does this metric enable?\n\nQ3: In Appendix E.1-3, you frequently use the approximation symbol \"\\approx.\" Is it possible to replace these approximations with inequalities or formal big-O notation to provide a more precise and formal explanation? How would this impact the rigor and clarity of the proofs presented?", "limitations": "A limitation of this study is that the model and data settings analyzed do not fully align with real-world transformers or language models.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The strength of this study lies in its theoretical analysis that connects specific roles, such as 'copier' and 'selector,' to each layer of the transformer when processing n-gram Markov chain data.", "weaknesses": "The analysis is specifically limited to n-gram Markov chain data, which may not reflect the complexity or diversity of data typically processed by language models. Furthermore, the structure of the feed-forward network is unusually specialized, and the learning process is divided into three distinct stages, which is not typical for standard transformer training. Additionally, the role of the attention mechanism is reduced merely to specifying positions within the n-gram model, diverging from its broader and more dynamic functionalities in conventional transformer applications. This specificity and simplification may limit the generalizability of the findings to more standard scenarios.\n\nAnother weakness of the study is the use of the Modified Chi-squared Mutual Information (MI) metric without sufficient explanation of its significance or relevance to real-world applications. \nWhile Chi-squared MI may theoretically is used in the total validation analysis,  it difficult for readers to grasp the practical implication to use the metric. Additionally, although  the following existing study analyze the learning of induction heads in bi-gram models, we need the adequate explaination of the fundamental differences between these methods and the approach taken in the study without the n-gram model. \nBirth of a Transformer: A Memory Viewpoint\nAlberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, Leon Bottou\nNeurIPS2023\n\n Furthermore, considering the short review periods typical of international conferences, the extensive proofs, despite its limited model, demand a high effort from reviewers.\n This adds significantly to the difficulties of this paper and makes it particularly laborious for reviewers.\n\nAdditionally, the proofs provided in Appendix E.1-3 of the study frequently use the approximation symbol \"\\approx,\" which introduces ambiguity into the explanations and raises concerns about the solidity and rigor of these proofs. This frequent reliance on approximations makes it challenging to determine whether the proofs are robust and reliable. This aspect is a significant drawback, as it undermines the confidence in the theoretical claims and the overall integrity of the research findings, casting doubt on the precision and applicability of the model’s theoretical underpinnings", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. Please see our response below.\n\n**On $n$-gram data, simplified model & split training phases**: \nWe remark that the $n$-gram Markov chain data assumption is much more expressive and challenging than the bi-gram data assumption in previous works (Bietti et al, (2023), Nichani et al, (2024)).\nRigorously analyzing the training dynamics for in-context learning of $n$-gram Markov chain is a significant open problem identified by Nichani et al, (2024), let alone with a more sophisticated transformer model that also include RPE, FFN, layer normalization, and residual link. \n\nIn this work, we address these challenges and reveal that the model learns a generalized induction-head (GIH) mechanism for $n$-gram data. We split the training phases to focus on the main hurdles. Experiments on simultaneous training of all parameters match our theoretical findings (see Figure 6 in Appendix C.1). Please refer to the general rebuttal for more details on model simplification and training phase splits.\n\n**On the modified $\\chi^2$-MI**: \nIn the generalized induction-head mechanism, the modified $\\chi^2$-MI criterion arises naturally from gradient calculations, not by design. For intuition, see Eqn (3.6) in Section 3.3, which relates the modified $\\chi^2$-MI to the standard $\\chi^2$-MI under certain symmetry conditions:\n$$\n\\log \\tilde I_{\\chi^2}(\\mathcal{S}) = \\log I_{\\chi^2}(\\mathcal{S}) - |\\mathcal{S}| \\log {d}. \\tag{1} \n$$\nThis criterion enables the model to balance maximizing mutual information and minimizing model complexity by selecting parents with high mutual information with the token to be predicted.\nFor instance, selecting a larger parent set in the GIH mechanism results in less biased prediction, while also incurring larger estimation variance. \nIn addition, we discuss in line 354-358 that the modified $\\chi^2$-MI will become equivalent to the standard $\\chi^2$-MI if the data is bi-gram, which reproduces the previous results on bi-gram data.\n\n\n**On the comparison to Bietti et al, 2023**: The major differences between our work and Bietti et al. (2023) are as follows:\n- *(Data structure)* Bietti et al. (2023) only consider 2-gram Markov data, while we consider general $n$-gram data. As discussed in the general rebuttals, going from bi-gram to $n$-gram data is highly nontrivial.\n- *(Model structure)* Bietti et al. (2023) do not consider an MLP layer or layer normalization and only use single-head attention for their construction of the induction head mechanism. In stark contrast, we keep all these key components and show they are necessary for learning $n$-gram data and exhibit a rich interplay. As our model is more complex, the original analysis in Bietti et al. (2023) does not apply.\n- *(Theoretical understanding)* Bietti et al. (2023) construct a hypothetical model with empirical evidence also on a synthetic dataset and provide an analysis for three gradient steps. Their analysis is essentially just around the initialization, which has no guarantee on the training trajectory or the limiting model. We theoretically understand the full training dynamics that lead to the emergence of the generalized induction head mechanism.\n\n**Can we \"scale up\" the analysis from bi-gram to $n$-gram?**\nSimply scaling up the bi-gram analysis in Bietti et al. (2023) does not work for $n$-gram data, as the original mechanism handles only bi-occurred tokens.\nLet us follow the construction in Bietti et al. (2023) and show a counterexample. \nConsider the simplest $3$-gram data, where each token depends on the previous two tokens as parents:\n- Let $x_l = w_E(l) + p_l$, where $w_E(l)$ is the word embedding for token $l$ and $p_l$ is the positional embedding for position $l$. Assume positional embeddings are orthogonal to each other, and $w_E(l)$ is orthogonal to all positional embeddings.  \n- The positional information is packed into the $W_{QK}$ matrix for the first layer, i.e., $W_{QK} = \\lambda \\sum_{l=2}^L \\sum_{s\\in \\{1, 2\\}} p_l p_{l-s}^\\top $ with $\\lambda$ being a large constant.\n- After the first layer's attention, the model will output \n$$\n    x_l' = W_{OV} x_{1:L}\\sigma(x_{1:l-1}^{\\top} W_{QK}^{\\top} x_{l}) \\approx W_{OV} (x_{l - 1} + x_{l -2}) = W_{OV} (w_E(l-1)+w_E(l-2) + p_{l-1} + p_{l-2}). \\tag{2}\n$$\nEven if this model can identify the parent sets, it fails to capture order information, i.e., distinguishing which parent is first and which is second. Therefore, without changing the model structure, the original construction in Bietti et al. (2023) fails for $n$-gram data.\n\n\n**Major challenges for $n$-gram data**: \nThe counterexample shows why single-head attention without additional nonlinearity fails on $n$-gram data. Our work leverages multi-head attention, FFN, and layer normalization to learn history-dependent features. The main challenge lies in analyzing the dynamics of a more sophisticated model with cross-head interference and nonlinearity, which Bietti et al. (2023) did not address. We will include these discussions in the revision.\n\n**On the proof and \"$\\approx$\" symbol**: \nWe want to clarify that Appendix E contains a proof sketch instead of a formal proof, and we use the \"$\\approx$\" symbol to facilitate understanding of the proof logic.\nThe complete and rigorous proofs are provided in Appendix F.\n\n\n**References**:\n- Bietti, Alberto, et al. Birth of a transformer: A memory viewpoint. 2023.\n- Nichani, Eshaan, et al. How transformers learn causal structure with gradient descent. 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2p2XAVq77e", "reviewer": "Reviewer_FHQ5", "summary": "In the paper, the authors analyze a simplified two-layer transformer model trained on mixtures of Markov chains, proving that gradient flow converges to a particular configuration of copier-selector-classifier network, that implements a generalized induction head mechanism. The authors further provide numerical experiments validating their findings.", "questions": "As mentioned above, I am a bit concerned about the removal of the word embedding layer and the fact that the first attention matrix is a function of just the positional embeddings. Is this simplification really necessary? What would break down in the proof with the introduction of word embeddings? Is there any hope to prove a similar result with word embeddings, and maybe with absolute embeddings, that to the best of my knowledge are used in all current state-of-the-art transformer models?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is well motivated and well written. The introduction of the generalized induction head estimator is original and effective. I checked the most important proofs and they look sound to me.", "weaknesses": "The study of full transformer models is certainly extremely complex, and the introduction of simplifications in the model is necessary for an effective theoretical study of the network. However, one simplification that might prevent an actual applicability of the authors' findings in real-world scenario is the complete removal of the word embedding layer, and the fact that the first attention matrix is a function of only the relative positional embeddings, without any information on the actual input tokens. Nonetheless, I believe that the study of the generalized induction head estimator is enough for the paper being worthy of publication in the conference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. Please see our response below.\n\n**First layer word embeddings**:\nAs stated in our general rebuttal, we study a simplified model for a better understanding of the functionality of each component in detail. \nAlthough the first attention layer does not contain trainable word embeddings, it can be considered as \n$$\n\\sigma(XWX + W_P)X,\n$$\nwhere we fix the word embeddings $W \\in \\mathbb{R}^{d \\times d}$ as a random matrix with Gaussian entries with mean $0$ and variance $1/d$. \nThis ensures that $x_l^\\top Wx_{l'} \\approx 0$ for all $l$ and $l'$ when the dimension $d$ is large. It can be interpreted as fixing $W$ at the initialization of the training, and it is common practice to set it as a random Gaussian matrix. Similar arguments can also be found in Bietti et al. (2023).\n\nWe remark that theoretical analysis of training dynamics becomes challenging when parameters for word embeddings are included. Previous works, such as Nichani et al. (2024) and Jelassi et al. (2022), often consider only positional embedding in the attention layer to focus on the positional information. Also, it has been empirically observed that the first attention layer's QK matrix has almost zero values for the token embedding parts during training (Nichani et al., 2024). This suggests that the word embeddings do not contribute to the induction head mechanism.\n\nFor our problem, when we incorporate trainable word embeddings, theoretical analysis of the training dynamics in Stage 2 becomes much more complicated. Since the first attention layer depends on both positional and word embedding weights, the dynamics of $w^{(h)}$ become intertwined with the dynamics of the word embedding weights. Then, simultaneous tracking of both dynamics is necessary, resulting in no closed-form solutions for these stochastic differential equations. A comprehensive analysis of the dynamics involving word embedding weights may be addressed in future work.\n\n**References**\n- Bietti, Alberto, et al. \"Birth of a transformer: A memory viewpoint.\" Advances in Neural Information Processing Systems 36 (2024)., \n- Nichani, Eshaan, Alex Damian, and Jason D. Lee. \"How transformers learn causal structure with gradient descent.\" arXiv preprint arXiv:2402.14735 (2024).,  \n- Jelassi, Samy, Michael Sander, and Yuanzhi Li. \"Vision transformers provably learn spatial structure.\" Advances in Neural Information Processing Systems 35 (2022): 37822-37836.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "89y24wrsv6", "reviewer": "Reviewer_ExgR", "summary": "This paper explores the theoretical foundations of in-context learning (ICL) within transformer architectures, particularly focusing on how various components contribute to ICL. The study examines a two-attention-layer transformer trained on n-gram Markov chain data, analyzing its role in ICL through gradient flow convergence to a limiting model that performs a generalized \"induction head\" mechanism. The key contributions include a rigorous theoretical analysis, highlighting the distinct roles of transformer components: the first attention layer as a copier, the feed-forward network with normalization as a selector, and the second attention layer as a classifier. The authors identify three phases of training dynamics and validate their theoretical findings with simulation experiments. This work provides a comprehensive understanding of how ICL is facilitated by the synergistic contribution of transformer components, bridging the gap between empirical observations and theoretical underpinnings, and offering valuable insights into transformer model design and training for complex tasks.", "questions": "- Can you provide empirical evidence or theoretical guarantees regarding the generalization ability of your model?\n\n- Do you plan to include additional experiments to validate your theoretical findings on more diverse datasets or more complex tasks?\n\n- In what sense do you use autoregressive transformers in your analysis?", "limitations": "The authors mention in the paper that Section C.1 contains a discussion of limitations, but this section is either missing or not clearly labeled.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The introduction of a generalized \"induction head\" mechanism and the detailed analysis of training dynamics in a two-attention-layer transformer trained on n-gram Markov chain data represent a creative combination of existing ideas, applied in a fresh and impactful way.\n\n- The theoretical analysis seems rigorous and well-supported by mathematical derivations and proofs. The authors provide a clear and logical progression from the problem setup to the main results, ensuring that each step is justified and well-explained. Though I did not check the appendix with details, the main content appears robust.\n\n- The paper is well-organized and clearly written, making complex concepts accessible. The use of diagrams and mathematical notation is appropriate and aids in the understanding of the key ideas.", "weaknesses": "- The work primarily focuses on the gradient flow of training dynamics. It would be beneficial to analyze gradient descent training to better align with practical implementations.\n\n- The paper could better contextualize its contributions by providing a more detailed comparison with prior work, such as [1]. This comparison should include aspects such as optimization methods, model structure, input embedding, and loss function. A comparison table is suggested to highlight differences and similarities clearly.\n\n- The paper makes several assumptions, particularly in the theoretical analysis. For example, it heavily relies on Relative Positional Embedding. Discussing the potential limitations these assumptions impose and how they might affect the generalizability of the results would be helpful.\n\n- The paper provides guarantees on training dynamics, but it would be beneficial to demonstrate the generalization ability of the proposed approach.\n\n- In line 318, there is a typo: it should be $\\partial_t a(t) \\asymp e^{a(t)}$.\n\n\n[1] Nichani, Eshaan, Alex Damian, and Jason D. Lee. \"How transformers learn causal structure with gradient descent.\" arXiv preprint arXiv:2402.14735 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. Please see our response below.\n\n**Gradient flow analysis**: We focus on gradient flow to simplify the theoretical analysis, as it is a common approach for understanding the training dynamics of complex models. \nWe believe that the results can be extended to gradient descent with a sufficiently small step size. \nGiven that our model includes two layers of multi-head softmax attention and feedforward neural networks, analyzing the training dynamics is challenging even with gradient flow. \nFurthermore, we use gradient descent in our experiments, and the results align with our theoretical findings. This demonstrates that our gradient flow analysis is consistent with the gradient descent implementation.\n\n**Comparison to Nichani et al. (2024)**:\nThe major distinctions between our work and Nichani et al. (2024) lie in the model structure and data structure. \nWe will include the following table for a clear comparison to Nichani et al. (2024). \n\n|                                          | Model structure                          | Causal data structure | Optimization methods | Input Embedding            | Loss function |\n|------------------------------------------|------------------------------------------|-----------------------|----------------------|----------------------------|---------------|\n| Our work                                 | Two layers of multi-head attention with FFN and layer normalization | Have $n$ parents      | Gradient flow        | RPE, Word embeddings | Cross entropy |\n| Nichani et al. (2024)                    | Two layers of single head attention      | At most one parent    | Gradient descent    | Positional Embedding, Word embeddings | Cross entropy |\n\n\n**Relative Positional Encoding**:\nOur findings highlight the benefits of using Relative Positional Encoding (RPE). Previous works have shown that transformers with RPE outperform those with standard positional encoding. For instance, Allen-Zhu et al. (2023) demonstrate that transformers with RPE are more effective at learning hierarchical language structures.\n\nIn our work, as we consider an $n$-gram data model, RPE is a suitable choice for positional encoding. Specifically, RPE has been utilized to effectively capture potential parent sets, thereby enhancing the length generalization ability of the model. As long as the model accurately identifies parent sets with RPE, it will also perform well when tested on input sequences of varying lengths $L$.\n\nWhile RPE may not be the best choice for different datasets, a generalization study of RPE is beyond the scope of our current work.\n\n**On Generalization ability**:\n\n- *Generalization to test data*:\n The generalization capabilities are tied to the model's ability to infer the true parent sets from the data. Our results demonstrate that the parameters are updated to maximize the modified chi-square mutual information. We find that the modified chi-square mutual information is an effective metric for identifying the parent sets as described in Section 3.3. For instance, when\n the number of parents is known a priori, the modified chi-square mutual information is maximized by the true parent set.\n\nWe have conducted additional experiments to evaluate the model's generalization performance on sequences having lengths different from the training data.\nPlease see the attached pdf for the experimental results under the general rebuttal.\n  These additional experiment results indicate that the trained model successfully captures the true parent sets, implying the strong performance on test data. Even if the test data have a different distribution or length \n $L$, the model will generalize well as long as they maintain the same parent structure as the training data. \n- *Generalization to other types of data/model*:\nThe investigation of data structures beyond the $n$-gram model and more general transformer model is beyond the scope of our current study.\nWe believe it is an important direction for future work to study the training dynamics of transformers in more general settings.\n\n\n**On Additional experiments**:\nThe goal of our paper is to examine the induction head mechanism of transformers within the $n$-gram data model framework. The $n$-gram model is widely used to evaluate transformers' capabilities in language modeling from both experimental and theoretical perspectives. Our experiments show that the empirical results align with our theoretical findings that the model implements an induction head. \nWe leave it as future work to explore the model's performance on more complex tasks and datasets.\n\n\n**Autoregressive Transformer in analysis**:\nWe want to clarify that our model is autoregressive in that it can predict the $(L+1)$-th token based on the previous $L$ tokens by outputing a distribution over the possible tokens. Importantly, $L$ can be arbitrarily large, provided it is sufficiently large to capture the context needed for accurate predictions.\n\n**Limitation**: We will include discussion on the limitations of our work in the revision.\n\n**Typo**: Thank you for pointing out the typo. We will correct it in the revision. \n\n**Reference**\nAllen-Zhu, Zeyuan, and Yuanzhi Li. \"Physics of language models: Part 1, context-free grammar.\" arXiv preprint arXiv:2305.13673 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# General Rebuttal to all Reviewers
We thank all reviewers for their valuable feedback. Below, we summarize our contributions and address common questions.

**Summary of contributions**:
To the best of our knowledge, our work is the first to show that through gradient-based optimization, a two-layer transformer can provably learn the induction-head mechanism on data from an $n$-gram Markov chain.
Our work features a rigorous analysis for the training dynamics of a two-layer transformer, organically incorporating key components including multi-head attention, relative positional embedding (RPE), feed-forward network (FFN), layer normalization, and residual link, while previous works only studied partial combination of these designs [Zhang et al. (2023); Huang et al. (2023); Bietti et al, (2023); Nichani et al, (2024)].

We reveal a generalized induction-head mechanism where the transformer selects a group of parents for inference based on the modified mutual information criterion. We demonstrate that this criterion naturally incurs a trade-off between model complexity and information richness.


**On the simplification**:
Given the complexity of transformer models and natural language, it is beneficial to study simplified models on synthetic data to understand the induction head mechanism while preserving the essence of real-world data and models.
In the following, we briefly discuss our assumptions on the data and the model:
- *(On $n$-gram data)* Natural language is generally considered Markovian, making the $n$-gram Markov chain a natural tool for language modeling (Coleman, 2005). 
Our current study significantly improves upon the previous theoretical works that are restricted to bi-gram Markov chains (Bietti et al, (2023), Nichani et al, (2024), Makkuva et al, (2024)).
    - In fact, going from bi-gram to $n$-gram is nontrivial as is discussed in Nichani et al, (2024), and our methodology is exploiting the multi-head attention mechanism with nonlinearity provided by FFN & layer normalization to learn a history-dependent feature. The rich interplay between these designs poses unique challenges for keeping track of the training dynamics.

- *(On transformer model)* We modify the original transformer model for better understanding of the functionality of each component in the transformer.
This is a common approach for theoretical studies (Bietti et al, (2024), Nichani et al, (2024)).
Nonetheless, we have kept the core components of the transformer, including multi-head attention, RPE, FFN, layer normalization, and residual link, to ensure that our findings are relevant to the real-world transformer models.
    - In fact, these modifications are based on the intuition provided by previous studies. For instance, the removal of the word similarity in the first attention's QK projection weights is inspired by the fact that the first layer's QK matrix has almost zero values for the token embedding parts during training (Nichani et al, (2024)) and does not contribute to the induction head mechanism.
    Rigorous study of the dynamics without simplifications like removing the token similarities is left for future work.

**On split training stages**: For a clearer theoretical understanding of the training dynamics, we split the training process into three phases. On one hand, this has a similar effect to using different learning rate for different layers, which is a standard practice in deep learning literature (Tian et al, (2023), Nichani et al, (2024), Lee et al, (2024)). On the other hand, the three-phase training paradigm can also be justified by the following facts:
- *(Theory perspective)* As is outlined between line 339-340, our results indicate a clear separation in the growth rates of different parameters in the model, which naturally gives rise to a multi-phase training dynamics where some weights like the FFN get trained first.
- *(Experiment perspective)* Experiments on simultaneously training all layers produce the same limiting model and the dynamics goes through a similar course, which also validates our theoretical findings. Please see Figure 6 in Appendix C.1 for more details.

Moreover, even with split of training phases, our dynamics analysis is still highly non-trivial given a more sophisticated model than the model studied in previous work.
In particular, we need to handle the cross-head interference and nonlinearity in FFN and layer normalization, which are not present in previous works (Zhang et al, (2023); Huang et al, (2023); Bietti et al, (2023); Nichani et al, (2024)).
The split of training phases helps us focus on the major hurdles in the dynamics analysis.
We will add more details, including the strengths and limitations in the revision when additional space is available.

**References**
- Zhang, Ruiqi, et al. Trained transformers learn linear models in-context. 2023.
- Huang, Yu, et al. In-context convergence of transformers. 2023.
- Coleman, John S. Introducing speech and language processing. 2005.
- Bietti, Alberto, et al. Birth of a transformer: A memory viewpoint. 2023.
- Nichani, Eshaan, et al. How transformers learn causal structure with gradient descent. 2024.
- Makkuva, Ashok Vardhan, et al. Attention with markov: A framework for principled analysis of transformers via markov chains. 2024.
- Tian, Yuandong, et al. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. 2023
- Lee, Jason D., et al. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. 2024.

Meta-review: The paper shows the learning of induction head through gradient dynamics analysis on Markov chain data in a mathematically rigorous manner. This is a novel and important contribution to the community, in which most of the papers studies its empirically behavior of induction head rather than building a theoretical framework. All reviewers agree the novelty of the paper and think the paper is well-written.
Decision: Accept (poster)