Title: Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: transformer

in context learning

Markov chain

n-gram model

feature learning

training dynamics

induction head
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

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. In the limiting model, 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.

1 Introduction

In-context learning (ICL) (Brown et al., 2020) has emerged as a crucial aspect of large language model (LLM) (Radford et al., 2019; Brown et al., 2020; Achiam et al., 2023; Anthropic, 2023; Team et al., 2023) functionality, enabling pre-trained LLMs to solve user-specified tasks during inference without updating model parameters. In ICL, a pre-trained LLM, typically a transformer, receives prompts containing a few demonstration examples sampled from a 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 of ICL requires the LLM to perform certain learning processes during inference.

Although many previous works aim to demystify ICL from either empirical or theoretical perspectives, 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 multihead 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 a gap in understanding how full-fledged transformer architectures facilitate ICL of more complex tasks, especially when latent causal structures exist among the tokens in a sequence.

In this paper, our aim is to narrow this gap by studying how a 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 n 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 a 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 -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-entropy loss converge during training? (ii) If 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 a copier , copying past tokens within a given window to each position. The FFN layer acts as a selector that generates a feature vector by only looking at informationally relevant parents from the window according to a modified χ 2 -mutual information. Finally, the second attention layer is 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 et 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 our best knowledge, our work is the first to 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 §D. Finally, we present corresponding experiment results in §B, and the detailed proofs are deferred to the Appendix.

Notation. We denote by e 1 , . . . , e d the standard basis vectors in R d and by 1 the all-one vector in R d . We denote by σ ( · ) the softmax function such that the i -th coordinate of σ ( x ) is σ i ( x ) =

exp( x i ) / ∑ L l =1 exp( x l ) for x ∈ R L . By default, the softmax operation will always be applied rowwise. For any integer n > 0 , we denote [ n ] := { 1 , . . . , n } . For a vector w ∈ R M , we denote by w i the i -th entry of w and w -i the ( M +1 -i ) -th entry of w for positive integer i ∈ [ M ] . For a matrix W , we denote by W ( i, j ) the entry at the i -th row and j -th column of W . For two vectors u and v , we write u/v as the vector obtained by taking element-wise division between u and v . We denote by a ∨ b and a ∧ b the maximum and minimum of a and b , respectively. We denote by x s : t the sequence { x s , x s +1 , . . . , x t } . For a class X , we denote by ∆( X ) 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. In 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 pretrained transformer can be viewed as a conditional distribution f tf ( · | 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 f tf is used to predict the ( L +1) -th token x L +1 with the first L tokens being the prompt. Here L is a fixed number and the joint distribution of the sequence x 1:( L +1) is sampled from a random n -gram Markov chain. In other words, with x 1:( L +1) sampled from some distribution, we evaluate how well f tf ( · | x 1: L ) predicts the distribution of x L +1 .

n -Gram Markov Chains. We assume the data comes from a mixture of n -gram Markov chain model, denoted by a tuple ( X , pa , P , µ 0 ) , where X is the state space and pa = ( -r 1 , . . . , -r n ) is the parent set with positive integers r 1 < r 2 < · · · < r n . That is, for each l > r n , x l only statistically depends on ( x l -r n , . . . , x l -r 1 ) , which is denoted by X pa ( l ) and referred to as the parent tokens of x l . We let

d = |X| denote the vocabulary size. Moreover, P is a probability distribution over the set of Markov transition kernels respecting the parent structure specified by pa , and µ 0 is the joint distribution of the first r n tokens x 1: r n . Note that the size of the parent set n can be smaller than or equal to r n . Thus, the sequence x 1:( L +1) is generated as follows: (i) sample initial r n tokens ( x 1 , . . . , x r n ) ∼ µ 0 , (ii) sample a random transition kernel π ∼ P , where π : X n → ∆( X ) , and (iii) sample token x l ∼ π ( · | X pa ( l ) ) for l = r n +1 , . . . , L +1 . See Figure 1 for an illustration of the generating model of x 1:( L +1) .

Cross-Entropy Loss. When x 1:( L +1) is generated, x 1: L is fed into the transformer f tf to predict x L +1 . To assess the performance of ICL, we adopt the population cross-entropy (CE) loss

where ϵ > 0 is a 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 x 1:( L +1) (including the randomness of π ∼ P ). When setting ϵ = 0 , we note that minimizing this crossentropy loss is equivalent to minimizing the KL divergence

As a remark, we also relax a condition in Nichani et al. (2024) where the last token x L has to be resampled from a uniform distribution. In addition, our analysis can also be extended to sequential CE loss, which corresponds to predicting every token in the sequence given the past rather than just the last token x L +1 . This is closer to the training paradigm used in practice (Brown et al., 2020). See §C.4 for a further discussion on the sequential CE loss.

2.2 A Two-Layer Transformer Model

We consider a class of two-attention-layer transformer model, denoted by TF ( M,H,d,D ) , which incorporates Relative Positional Embedding (RPE) (He et al., 2020), Multi-Head Attention (MHA) (Vaswani et al., 2017), and a Feed-Forward network (FFN) with normalization. Here M is an integer that specifies the window size of RPE, H is the number of heads in the first attention layer, d is the vocabulary size, and D is an integer that controls the complexity of FFN. The details of TF ( M,H,d,D ) are as follows.

Token Embedding, Input and Output. Note that each token takes values in X with d = |X| . We embed the tokens into one-hot vectors in R d , and thus we can identify X as the canonical basis in R d , i.e., X = { e 1 , . . . , e d } . A transformer model can be viewed as a mapping from R ( L +1) × d to ∆( X ) . In particular, given the input sequence x 1: L , we denote X = ( x 1 , . . . , x L ) ⊤ ∈ R L × d , and we append a zero vector 0 ∈ R d to the sequence, and define r X = ( x 1 , . . . , x L , 0 ) ⊤ ∈ R ( L +1) × d . The transformer takes r X as input and outputs a probability distribution over X .

𝑊 𝑃

(ℎ)

𝑤 -1

ℎ

𝑤 -2

ℎ

…

𝑤 -𝑀

ℎ

𝑤 (ℎ)

-∞

(a)

(b)

Relative Positional Embedding. In each head of the first attention layer, we adopt RPE to incorporate positional information. Specifically, RPE is parameterized by a vector w = ( w -M , . . . , w -1 ) ⊤ ∈ R M , and it assigns a scalar value W P ( i, j ) to a pair of positions ( i, j ) satisfying

In other words, as illustrated in Figure 2, the i -th token only attends to tokens with indices in { i -1 , . . . , i -M } , referred to as the lengthM 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 position.

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 r X into a sequence in R d with length L +1 , denoted by V ( h ) = ( v ( h ) 1 , . . . , v ( h ) L +1 ) ⊤ ∈ R ( L +1) × d . For any l ∈ [ L +1] , v ( h ) l is computed via

That is, we use the RPE parameter W ( h ) P to construct a weighted sum over the input sequence at each position l ∈ [ L +1] . Here W ( h ) P 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 (1) , . . . , V ( H ) ) ∈ R ( L +1) × Hd . Here we abuse the notation and write V = ( v 1 , . . . , v L +1 ) ⊤ , i.e., each v l is the l -th row of V . For any vector v ∈ R Hd , we can split it into ( v (1) ⊤ , . . . , v ( H ) ⊤ ) ⊤ where each block v ( h ) ∈ R d . For embedding dimension d e , each vector of V is passed through an FFN ϕ ( · ) : R Hd → R d e , which specifies a polynomial kernel such that for any v, v ′ ∈ R Hd , we have

Here, the low-degree parent set [ H ] ≤ D := {S ⊆ [ H ] : |S| ≤ D } contains all subsets of [ H ] with cardinality at most D , and { c S : S ∈ [ H ] ≤ D } are the corresponding trainable parameters of ϕ ( · ) . Therefore, the FFN ϕ ( · ) specifies a kernel on the output of the multihead attention which induces a special inner product structure. While (2.3) characterizes ϕ ( · ) implicitly, we provide an explicit construction of ϕ ( · ) in Lemma C.1 as a vector-valued mapping whose entries are monomials of the

input's entries. Moreover, the complexity of ϕ ( · ) is controlled by the maximum degree D , which also influences the embedding dimension d e as we show in the construction.

Furthermore, to control the magnitude of the FFN outputs, we normalize ϕ ( · ) by letting u l = ϕ ( v l ) / √ C D for all l ∈ [ L + 1] , where we define C D = ∑ S∈ [ H ] ≤ D c 2 S . Such a normalization scheme is motivated by the standard layer normalization (Ba et al., 2016) in transformer architectures. To motivate the use of √ C D as the normalization, consider a special case where the positional embeddings, after the softmax function, produce attention weights that are close to one-hot for each head. Then v ( h ) l in (2.2) is equal to some token in x 1: L . As a result, each v l consists of H tokens and

Thus, u l is roughly equivalent to the output of the layer normalization ϕ ( v l ) / ∥ ϕ ( v l ) ∥ 2 (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 1 , . . . , u L +1 ) ⊤ and the original sequence r X are then fed into the second attention layer to generate the final output. In particular, u L +1 is used as the query to compare with the keys { u M +1 , . . . , u L } , and the resulting attention scores are used to aggregate the values x ( M +1): L . This attention layer has a single head and a scalar trainable parameter a . We let U 1: L = ( u 1 , . . . , u L ) ⊤ ∈ R L × d e and denote by Mask ( · ) the mask that sets every entry of the first M rows of a matrix to be -∞ . The final output is given by

Note that the softmax function in (2.4) yields a probability distribution over [ L ] and that x 1: L is a sequence of one-hot vectors. Thus y in (2.4) is a probability distribution over X . The mask operator is included here just to simplify our analysis while in the experiments we are not using the mask.

In summary, given the input r X ∈ R ( L +1) × d , in the matrix form, our transformer model TF ( M,H,d,D ) consecutively applies the following operations:

The trainable parameters of the above transformer model are denoted by

We remark that the transformer model in (2.5) is known as a disentangled transformer (Friedman et al., 2024), which is a version of the transformer model that is more amenable for theoretical analysis. One thing to be noted is that there is a residual connection that directly copies r X to the output of the FFN & Normalize block, which gives us [ U, r X ] , and the second attention layer will treat the copied r 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.

Our goal is to investigate whether the transformer model TF ( M,H,d,D ) can perform ICL over n -gram Markov chains and further, whether such capability can be learned from data with common training algorithms like gradient descent.

3 Theoretical Results

In this section, we present the theoretical results. We first show in §3.1 and §C.1 that there exists a transformer in TF ( M,H,d,D ) that implements a generalized 'induction head' mechanism (Olsson et al., 2022) with a learned feature, which serves as a natural algorithm for learning n -gram Markov chains. Then in §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 , pa , P , µ 0 ) 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 π ∈ supp( P ) , where a rigorous statement is deferred to Assumption 3.5. We also assume the window size M > r n . For any n -gram Markov chain with transition kernel π ∼ P , we let µ π ∈ ∆( X M +1 ) denote the stationary distribution of the Markov chain over a window of size M +1 . Here we use { z ℓ } l ≥ 1 to denote a random sequence of tokens generated by the Markov chain. Then µ π denotes the joint distribution of a block of M +1 tokens ( z l -M , . . . , z l -1 , z l ) under the stationary distribution of π , where l > M is an integer.

In the following, we introduce a generalized induction head (GIH) estimator for the task of predicting x L +1 given x 1: L , which is based on the following simple idea: x L +1 should be similar to a previous token x l if their parents are similar . As the parent set pa is unknown, GIH adopts an informationtheoretic criterion to select a subset of previous tokens as a proxy of the parents. Specifically, GIH uses a modified version of χ 2 -mutual information, which is defined as follows.

Definition 3.1 (Modified χ 2 -Mutual Information) . We take a length-( M + 1) windows ( z l -M , . . . , z l -1 , z l ) for some l > M and suppose the sequence is sampled from stationary distribution µ π with π ∼ P . Let Z = ( z l -M , . . . , z l -1 ) . For any subset S ⊆ [ M ] , we use Z -S to denote the subvector of Z containing entries of the form z l -s , ∀ s ∈ S . For instance, suppose S = { 2 , 5 } , then Z -S = ( z l -5 , z l -2 ) . The modified χ 2 -mutual information for S is defined as

where µ π ( z = · | Z -S ) is the conditional distribution of z induced by µ π given the partial history Z -S , and µ π ( Z -S ) , µ π ( z ) are the marginal distributions of Z -S and z under ( z, Z ) ∼ µ π .

Intuitively, r I χ 2 ( S ) is modified from the vanilla χ 2 -mutual information ( χ 2 -MI) between two random variables (Polyanskiy and Wu, 2024) and quantifies how much information the partial history Z -S contains about z . In particular, we incorporate an additional µ π ( Z -S ) term that decreases with the growing size of S . To see the rationality, we first introduce a GIH estimator based on the modified χ 2 -mutual information.

Definition 3.2 (Generalized Induction Head) . A GIH estimator with window size M ∈ N , feature size D ∈ N is denoted by GIH ( · ; M,D ) , which maps x 1: L to a distribution over X . We let S ⋆ be the information-optimal subset (referred to as the 'information set' in the sequel 2 ) of [ M ] with size no more than D that maximizes the modified χ 2 -mutual information r I χ 2 ( · ) defined in (3.1) . That is, we define the information set S ∗ as

Then GIH ( x 1: L ; M,D ) outputs

Here, we define X l -S ⋆ as the set { x l -s : s ∈ S ⋆ } and N = ∑ L l = M +1 1 ( X l -S ⋆ = X L +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 r I χ 2 ( · ) defined in (3.1). In a nutshell, the GIH estimator checks

whether the partial histories of X l -S ⋆ and X L +1 -S ⋆ match and aggregate all the tokens x l that have a matching partial history to predict x L +1 . As a remark, using the modified χ 2 -MI 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, z and Z -S are conditionally independent given the true parents Z pa . Thus, maximizing the vanilla χ 2 -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 a strictly smaller r I χ 2 compared to the exact parent set, making them suboptimal.

The modified χ 2 -MI selects informative partial history. Even a 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 to appear less frequently in the context sequence, leading to poor estimation accuracy for the estimator in (3.3). The modified χ 2 -MI reaches a balance by selecting the informative partial history while penalizing the size of the information set.

The term involving µ π ( z = · | Z -S ) 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 χ 2 -MI strikes a balance between these two objectives and enables us to find a good proxy S ⋆ of pa when L is finite. Moreover, when L is sufficiently large, we identify two scenarios in which maximizing r I χ 2 ( · ) yields the true parent set (see §C.7 for details). Moreover, the GIH estimator is a 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 §C.1, the GIH mechanism can be implemented by the transformer model.

3.2 Convergence Guarantee of Gradient Flow

In the following, we present the convergence guarantee for gradient flow. To simplify the discussion, we consider the case where H = M , meaning there are enough heads to implement the GIH mechanism by having each head copy a unique parent token from a window of size M . Let us first introduce the paradigm of training by gradient flow.

Training Paradigm. Consider training a transformer TF ( M,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 L (Θ) as the population cross-entropy loss in (2.1), where the transformer model f tf is given by (2.5) with a parameter Θ . Ideally, when training the parameter Θ with gradient flow, the dynamics with respect to the loss L (Θ) is given by:

We consider a three-stage training paradigm where, in each stage, only a specific subset of the weights is trained by gradient flow. The three stages are outlined in Table 1. 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 a 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 a 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 [ H ] ≤ D as

where we recall that S ⋆ defined in (3.2) maximizes the modified χ 2 mutual information.

Assumption 3.3 (Initialization) . We assume that the following holds at initialization:

̸

For the first attention layer's RPE weights, w ( h ) -h ≥ w ( h ) -j +∆ w for all h, j ∈ [ H ] with j = h , where ∆ w > 0 is a positive scalar satisfying

The scalar parameter a in the second attention layer satisfies 0 < a ≤ O ( L -3 / 2 ) .

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 ICL.

We now outline our assumptions on the Markov chain used in the data generation process. Recall that r n is the largest absolute integer in the parent set pa . For any position l , we define the history Z = ( z l -r n , . . . , z l -1 ) as the last state and Z ′ = ( z l -r n +1 , . . . , z l ) as the current state. Since the parent of the new token z l is already included in Z , Z ′ is independent of all prior history given Z , forming a Markov chain.

We define P π as the d r n × d r n transition matrix for this Markov chain, where states are successive r n -tokens. Each row of P π is indexed by Z ′ and each column by Z . The matrix element P π ( Z ′ , Z ) is thus given by

This means that to transition from Z to Z ′ , all elements of Z ′ except for z ′ -1 must match the last r n -1 tokens of Z . The token z ′ l is then sampled according to the transition kernel π and depends only on the parent Z pa ( l ) . The above definition is in fact independent of the position l 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 P π .

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 P k are positive.

We defer more details about the above definition to §C.3. By the celebrated Perron-Frobenius theorem, if 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 is the stationary distribution. Next, we state the assumptions on the mixture of Markov chains for data generation.

Assumption 3.5 (Markov Chain) . For any π ∈ supp( P ) , we assume that:

The transition matrix P π is primitive. In particular, we assume that there exists λ < 1 such that the eigenvalue of P π with the second largest magnitude satisfies | λ 2 ( P π ) | ≤ λ . Note that λ 2 ( P π ) can be complex-valued.

There exists γ > 0 such that the transition kernel satisfies π ( x | X pa ) ≥ γ for any ( x, X pa ) .

In fact, the second condition π ( · | X pa ) > γ already ensures that P π must be primitive, as is required by the first condition. See Corollary F.14 for details. On the high level, the first assumption guarantees a unique stationary distribution as well as a fast mixing rate of the Markov chain by ensuring a spectral gap for P π . The second assumption implies a lower bound on the probability for any set S ⊆ [ M ] under the stationary distribution, i.e., µ π ( X l -S ) ≥ γ |S| for any l > M . See Corollary F.15 for details.

Now we are ready to present our main theoretical result on training transformers by gradient flow.

Theorem 3.6 (Convergence of Gradient Flow) . Suppose Assumption 3.3 and Assumption 3.5 hold. Consider H ≥ M . We set ε = L -1 / 2 for the cross-entropy loss and assume L is sufficiently large. Then the following holds for the three-stage training of gradient flow:

Then in the first stage with duration t 1 ≍ C D (0) log L/ ( a (0)∆ r I χ 2 ) , the ratio c S ⋆ /c S grows exponentially fast for any S ̸ = S ⋆ , and S ⋆ dominates exponentially fast in the sense that,

Stage II: Concentration of The First Attention. Define σ ( h ) ( t ) = σ ( w ( h ) ( t )) ∈ R M , and let σ min ( t ) := min h ∈S ⋆ σ ( h ) -h ( t ) . Then in the second stage with duration t 2 -t 1 ≍ L/ ( a (0)∆ r I χ 2 ) , the first layer's attention heads have attention probabilities concentrated on the optimal information set S ⋆ in the sense that for any t ∈ [ t 1 , t 1 + t 2 ) ,

Stage III: Growth of The Second Attention. For some constants c 1 , c 2 depending on ( P , S ⋆ ) with 0 < c 1 < c 2 , there exists a small constant δ > 0 such that the growth of a ( t ) exhibits the following two sub-stages: (i) When a ( t ) ≤ log( c 1 /δ ) , it holds that ∂a ( t ) ≍ e a ( t ) ; (ii) After a ( t ) has grown such that a ( t ) ≥ log( c 2 /δ ) , then ∂ t a ( t ) ≍ 1 /a ( t ) until it reaches the value log L/ 8 .

See §D for a proof sketch and §E for the detailed proof. We require that L is sufficiently large, and the specific conditions for L are deferred to §E.1.

Interpretation of Training Dynamics. We empirically verify Theorem 3.6 by conducting a simulation experiment. In particular, we train a transformer with H = 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 a 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 p S ⋆ ( t ) ≈ 1 and p S ( t ) ≈ 0 for all S ̸ = S ⋆ . In this case, for any input vectors v, v ′ ∈ R Hd , the inner product in (2.3) reduces to

That is, FFN only selects the blocks in S ⋆ 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 c S ⋆ immediately dominates the rest of c S 's 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 a 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 ∈ S ⋆ , σ ( h ) ( t ) = σ ( w ( h ) ( t )) ≈ 1 . Recalling the construction of RPE, this implies that v ( h ) l in (2.2) becomes x l -h for all

D=2, train separately

Stage III: Train

𝑎

Stage II: Train

𝑤 (ℎ)

Stage I: Train

𝑐 𝒮

(a)

(b)

(c)

Beginning

Ending

Epoch

Epoch

h ∈ 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 p S ⋆ ≈ 1 .

After the first two stages are completed, we know that the features constructed approximately satisfy (C.1) up to a proportionality factor. Then, in the final training stage, the scalar weight a in the second attention layer keeps increasing. Thus, this stage learns an exponential kernel classifier as specified in (C.2). When a ( t ) is sufficiently large, the learned transformer is close to a classifier that uses covariate-label pairs of the form ( X l -S ⋆ , x l ) to predict x L +1 . In particular, when a ( t ) goes to infinity, the transformer exactly becomes the GIH mechanism given in Definition 3.2. Moreover, we theoretically prove that the increasing trajectory of a ( t ) has two stages, where d a ( t ) / d t 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 approximately implements the GIH mechanism. We will prove that the difference between these two estimators is at most O ( L -1 / 8 ) . We defer the formal statement and proof to §E.5. Moreover, as an answer to the Question (iii) raised in §1, the different components of the transformer architecture are 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 a 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 to §B.2.

In §C.7, we provide more intuitive interpretation of the modified χ 2 -mutual 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 a two-attention-layer transformer model for learning n -gram Markov chains in an in-context way. Our work opens new directions for developing a 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.

References

Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S. et al. (2023). Gpt-4 technical report. arXiv preprint arXiv:2303.08774 .

Ahn, K., Cheng, X., Daneshmand, H. and Sra, S. (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.00297 .

Ahuja, K., Panwar, M. and Goyal, N. (2023). In-context learning through the bayesian prism. arXiv preprint arXiv:2306.04891 .

Akyürek, E., Schuurmans, D., Andreas, J., Ma, T. and Zhou, D. (2023). What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations .

Alayrac, J.-B., Donahue, J., Luc, P., Miech, A., Barr, I., 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 23716-23736.

Anthropic (2023). Model card and evaluations for claude models.

Ba, J. L., Kiros, J. R. and Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv:1607.06450 .

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 arXiv:2306.04637 .

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, F. and Kempe, J. (2024). Iteration head: A mechanistic study of chain-of-thought. arXiv preprint arXiv:2406.02128 .

Chen, S. and Li, Y. (2024). Provably learning a multi-head attention layer. arXiv preprint arXiv:2402.04084 .

Chen, S., Sheen, H., Wang, T. and Yang, Z. (2024). Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442 .

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, S. (2023). Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528 .

Collins, L., Parulekar, A., Mokhtari, A., Sanghavi, S. and Shakkottai, S. (2024). In-context learning with transformers: Softmax attention adapts to function lipschitzness. arXiv preprint arXiv:2402.11639 .

Deora, P., Ghaderi, R., Taheri, H. and Thrampoulidis, C. (2023). On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.12680 .

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.11004 .

Edelman, B. L., Goel, S., Kakade, S. and Zhang, C. (2022). Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning . PMLR.

Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T. 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.17086 .

Giannou, A., Rajput, S., Sohn, J.-Y., Lee, K., Lee, J. D. and Papailiopoulos, D. (2023). Looped transformers as programmable computers. In Proceedings of the 40th International Conference on Machine Learning (A. Krause, E. Brunskill, K. 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 transformers emulate in-context newton's method? arXiv preprint arXiv:2403.03183 .

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 arXiv:2310.10616 .

He, J., Chen, S., Zhang, F. and Yang, Z. (2024). From words to actions: Unveiling the theoretical underpinnings of llm-driven autonomous systems. arXiv preprint arXiv:2405.19883 .

He, P., Liu, X., Gao, J. and Chen, W. (2020). Deberta: Decoding-enhanced bert with disentangled attention. arXiv preprint arXiv:2006.03654 .

Honovich, O., Shaham, U., Bowman, S. R. and Levy, O. (2022). Instruction induction: From few examples to natural language task descriptions. arXiv preprint arXiv:2205.10782 .

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. Advances in Neural Information Processing Systems , 35 37822-37836.

Jeon, H. J., Lee, J. D., Lei, Q. and Van Roy, B. (2024). An information-theoretic analysis of incontext learning. arXiv preprint arXiv:2401.15530 .

Kim, J. and Suzuki, T. (2024). Transformers learn nonlinear features in context: Nonconvex meanfield dynamics on the attention landscape. arXiv preprint arXiv:2402.01258 .

Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. S. and Oymak, S. (2024). Mechanics of next token prediction 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 arXiv:2303.04245 .

Lin, L., Bai, Y . and Mei, S. (2023). Transformers as decision makers: Provable in-context reinforcement 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.10749 .

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.03576 .

Makkuva, A. 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.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 arXiv:2402.04161 .

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.10510 .

Nichani, E., Damian, A. and Lee, J. D. (2024). How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 .

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. OpenAI blog , 1 9.

Rajaraman, N., Bondaschi, M., Ramchandran, K., Gastpar, M. and Makkuva, A. V. (2024a). Transformers on markov data: Constant depth suffices. arXiv preprint arXiv:2407.17686 .

Rajaraman, N., Jiao, J. and Ramchandran, K. (2024b). Toward a theory of tokenization in llms. arXiv preprint arXiv:2404.08335 .

Sanford, C., Hsu, D. and Telgarsky, M. (2023). Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.02896 .

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, J. and Zhong, Y. (2023). Uncovering hidden geometry in transformers via disentangling position and context. arXiv preprint arXiv:2310.04861 .

Tarzanagh, D. A., Li, Y., Thrampoulidis, C. and Oymak, S. (2023a). Transformers as support vector machines. ArXiv , abs/2308.16898 .

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 models. arXiv preprint arXiv:2312.11805 .

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 1-layer transformer. arXiv preprint arXiv:2305.16380 .

Tian, Y., Wang, Y., Zhang, Z., Chen, B. and Du, S. (2023b). Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535 .

Vasudeva, B., Deora, P. and Thrampoulidis, C. (2024). Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738 .

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. and Polosukhin, I. (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: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593 .

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, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D. et al. (2022). Chain-of-thought prompting elicits reasoning in 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.08391 .

Xie, S. M., Raghunathan, A., Liang, P. and Ma, T. (2021). An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 .

Zhang, R., Frei, S. and Bartlett, P. L. (2023a). Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 .

Zhang, Y., Liu, B., Cai, Q., Wang, L. and Wang, Z. (2022). An analysis of 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 preprint arXiv: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 arXiv:2205.10625 .

Contents

Organization of The Appendix

The appendices are organized as follows:

In §A, we present an in-depth discussion on 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 in the main text.

In §D, we provide a high-level overview of the proof of our main results.

In §E, we present the proof for Theorem 3.6.

In §F, we 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 in 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-modalities (Alayrac et al., 2022). Moreover, recent 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 et al., 2022; Sanford et al., 2023; Giannou et al., 2023; Liu et 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).

There is 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 et al. (2021); Muller et al. (2021); Zhang et al. (2022, 2023b); Ahuja et al. (2023); Jeon et al. (2024); He et al. (2024) and the references therein. Another line of work investigates how transformers internally emulate specific algorithms to solve ICL tasks, where Akyürek et al. (2023); Von Oswald et al. (2023); Fu et al. (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 et al. (2023); Collins et al. (2024); Guo et al. (2023) investigate transformers' capabilities in learning with nonlinear functions. However, all of these works above focus on 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 et al., 2024; Makkuva et al., 2024b; Chen

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 different 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 investigated under different data models for both regression and classification tasks (Zhang et al., 2023a; Huang et 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 two-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 a surge of interest in understanding the induction head mechanism and its role in 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 a 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 -gram 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 et al. (2024a) also prove that training a single layer attention with a 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 a 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 to extend the analysis to more general n -gram Markov chains.

In this work, we show that a generalized version of the induction head mechanism can emerge when training a two-layer transformer on n -gram Markov chains. Moreover, our transformer models are more sophisticated, incorporating features like relative positional embedding, multi-head attention, 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 W Q , W K , W V and the output embedding matrix W O in the first attention layer. Let us 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 } as the parent set. Thus, the number of parents is n = 2 and the token embedding dimension is d = 3 . Note that for each sequence, the transition matrix π ( x | x pa ) is of shape d × d n . We assign a 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 α = 0 . 01 , i.e., π ( ·| x pa ) ∼ Dir ( α · 1 d ) . 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.

D=2, train separately

Stage III: Train

𝑎

Stage II: Train

𝑤 (ℎ)

Stage I: Train

𝑐 𝒮

(a)

(b)

(c)

Beginning

Ending

Epoch

Epoch

B.1 Training with Stage Splitting

we present the simulation results with model TF ( M,H,d,D ) in (2.5) and training in the three-stage manner. We configure the model with window size M = 3 , number of heads H = 3 , vocabulary size d = 3 and maximal FFN degree D = 2 .

Model initialization. The RPE weight matrix W ( h ) P is initialized such that the ( -i ) -th diagonal of W ( h ) P has value w ( h ) -i for i = 1 , 2 , . . . , M , while all other entries are initialized to -∞ . See Figure 2 for an interpretation. We initialize w ( h ) -h = 3 and set the remaining entries within the sizeM 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, all c S for S ∈ [ 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 a constant learning rate that is set to one for all stages. We train the model in Stage I (update parameters { c S } only) for 2000 epochs, in Stage II (update parameters { w ( h ) } only) for 50,000 epochs, and in Stage III (update parameter a only) for 5000 epochs, respectively. All experiments are conducted using a single Nvidia A100 GPU. The results are shown in Figure 4, which matches our theoretical results.

0.000100

train_loss

1.0-

0.000095-

0.8

000

001

0.000090

9:0

010

011

0.000085

0.4

100

101

0.000080

0.2

110

0.000075

0.0

0

2000

4000

6000

800010000

0

2000 400060008000 10000

Epoch

Epoch

(a)

(b)

0

0

1.0

1.0

0.6

2

2

0.8

0.8

0.5

4

0.6

4

0.6

0.4

6

0.4

6

0.4

0.3

0.2

0.2

0.2

8

8

0.1

0.0

0.0

0

i

2

0

1

2

0.0

0

2000

4000

6000

8000

10000

Beginning

Ending

Epoch

(c)

(d)

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.

In head h of the first attention layer, W ( h ) P is the relative positional embedding matrix, and we include W ( h ) Q ∈ R d × d , W ( h ) K ∈ R d × d and W ( h ) V ∈ R d × d as the weight matrices for the query, key, and value

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 r X to the output of the first attention layer. For the FFN layer, ϕ : R ( H +1) d → R d e is the same feed-forward network specified in (2.3). Here, we use a standard ℓ 2 -layer-normalization LN ( · ) , defined as

The second attention layer takes X as the value, which comes from the residual link (i.e., concatenation of r U and V while r X in V remains the same after ℓ 2 -normalization). In comparison to the simplified model in (2.5), here we incorporate the query, key and value projections for the first layer as in a 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 W ( h ) , W ( h ) and W ( h ) are initialized as identity matrices scaled by 0.001.

Q K V We initialized the RPE vector w ( h ) as w ( h ) -h = 1 for h = 1 , 2 , 3 , and leave the remaining entries within the lengthM 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 5, the full model converged to a state comparable to our simplified model. We further plot the W (1) Q , W (1) K , W (1) V for the first head after training in Figure 6. 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 ) V to be identity matrix.

W(1)

Q

0.5

0.0100

-0.5

0.0100

0.5

0.0

0.0075

0.0

0.0075

0.0

0.06

0.0050

0.0050

0.5

0.0025

0.5

0.0025

0.5

0.04

1.0

0.0000

1.0

0.0000

1.0 -

0.02

1.5

-0.0025

1.5

-0.0025

1.5

0.00

-0.0050

-0.0050

2.0

-0.0075

2.0

-0.0075

2.0

-0.02

2.5

-0.0100

2.5

-0.0100

2.5

0.04

0

1

2

0

1

2

0

1

2

(a)

(b)

(c)

B.3 Prior and Length Generalization

We further test the model learned by the three-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 §3.1, the GIH estimator can be applied to a sequence with an arbitrary length and does not concern the prior distribution of the underlying Markov chain. Thus, it is natural to see if the pre-trained transformer can also generalize to different lengths and prior distributions.

Recall that we train the transformer model with sequence length L = 100 and the concentration parameter of the Dirichlet prior is α = 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

concentration parameter α , we sample a random Markov chain, and generate a sequence of length L , and evaluate of cross-entropy loss for predicting x L +1 . Here we choose α ∈ { 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 are shown in Figure 7. The results show a decreasing trend in testing loss as the sequence length increases. For α = 0 . 2 , we observe first a small increase in the test loss when L just exceeds 100 , but then the loss decreases as L increases further. This experiment shows that the pre-trained transformer indeed generalizes in length and is robust to the change of prior distribution.

le-4

8.0

α=0.05

α =0.1

Validationloss

7.5

α=0.2

7.0

6.5

6.0

101

102

103

Sequencelength

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 to this transformer in the limit.

XL+1

x么

Value

FFN &LN

* = {1,2}

Key

Query

Head 3

X1-3

Head 2

X1-2

Residual

x么

0

UL+1

Head 1

x1-1

X1-3

X1-2

X1-1

x1

x1

0

head h ∈ S ⋆ can attend to the h -th parent token by setting the RPE weights in the softmax function to be w ( h ) = ρ · e -h for a sufficiently large ρ , where e -h ∈ R M is the canonical basis vector with the ( M +1 -h ) -th entry being one and all other entries being zero. As a result, each v ( h ) l for h ∈ S ⋆ satisfies v ( h ) l ≈ x l -h .

Step II: FFN Generates the Polynomial Features of the Information Set S ⋆ . As we have introduced in (2.3), each learnable c S 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 c S ⋆ = 1 whereas c S = 0 for all S ̸ = S ⋆ . The exact form of the output of the FFN layer, ϕ ( v l ) , is deferred to §C.2. Here the only property we require is that

Here X l -S ⋆ := ( x l -s : s ∈ 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 s l in (C.1) as an indicator for whether the information set of a token x l matches the information set of the token x L +1 . Then for the second attention layer, by setting a to be sufficiently large, the output will become

where N = ∑ L l = M +1 1 ( X l -S ⋆ = X L +1 -S ⋆ ) . That is, if at least one token x l has a matching information set as x L +1 , i.e., their histories restricted to S ⋆ are the same, the second attention layer outputs the average of such tokens. Otherwise, it outputs the average of previous tokens from x M +1 to x L . In Lemma E.6 in the appendix, we will show that the model learned by gradient flow implements the GIH mechanism up to a diminishing approximation error.

The weights of the transformer constructed above are illustrated in Figure 9. We consider the transformer model with M = H = 3 , d = 3 , and D = 2 . In this case, in the first attention layer, for each h ∈ [3] , W ( h ) P has three finite parameters w ( h ) -1 , w ( h ) -2 , and w ( h ) -3 . By our construction, we have w ( h ) -h = ρ for all h ∈ [3] and the rest of the entries of { w ( h ) } h ∈ [3] are all equal to zero. In Figure 9-(a) we plot the top ten by ten block of W (1) P , where w (1) -1 = ρ is shown in yellow and w (1) -2 = w (1) -3 are shown in purple. The gray color stands for -∞ entries. In Figure 9-(b) we plot { w ( h ) } h ∈ [3] . 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 , 3 } , 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 } , c S ⋆ = 1 , and c S = 0 for any other S .

C.2 Feed-Forward Network for Polynomial Kernel

Lemma C.1. Recall the FFN satisfying (2.3) , which maps a vector z ∈ R dH to a vector in R d e . We write z as ( z (1) , . . . , z ( H ) ) where z ( h ) ∈ R d for all h ∈ [ H ] . Let z ( h ) i be the i -th entry of z ( h ) . Then we can explicitly construct ϕ ( · ) by letting

which is equivalent to

where vec( · ) is the vectorization operator that transforms a tensor into a vector by stacking all the entries in the tensor. That is, for any S , we consider the |S| vectors in R d , { z ( h ) } h ∈S . In (C.3) we

(a)

(b)

(c)

𝑤

(")

𝑊

$

(")

Gray represents

-∞

𝑐

𝒮

𝒮

𝑐

𝒮

⋆

dominates

𝑤

(&)

𝑤

(')

0

8

0

8

1.0

7

1

7

2

6

2

6

0.8

5

5

0.6

4

4

4

6

3

3

0.4

2

2

0.2

8

1

1

0

0

0.0

0

2

4

6

8

0

1

2

000

001

010

011

100

101

110

compute all possible products of the entries of these vectors and multiply them by c S . In particular, for each S ∈ [ H ] ≤ D , we enumerate i h ∈ [ d ] for all h ∈ S . Therefore, the output dimension of ϕ is given by

Proof. First, we note that the indices of ϕ ( · ) have a grouped structure - we first enumerate all subsets in [ H ] ≤ D and then enumerate all monomials with superscripts in S . Since there are d |S| monomials, the output dimension is given by (C.4).

It remains to verify (2.3) with ϕ ( · ) defined in (C.3). To this end, we note that for any u, v ∈ R dH and any S ∈ [ H ] ≤ D , we have

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 Perron-Frobenius theorem on non-negative matrices (Meyer, 2023, Chapter 7). We consider the following class of irreducible matrices.

̸

Definition C.2 (Irreducible Matrix) . A non-negative square matrix P ∈ R d × d + is called irreducible if the induced directed graph G is strongly connected, i.e., for any pair of nodes in the graph, there always exists a directed path that connects these two nodes. Here, the induced graph G is defined based on d nodes with adjacent matrix A given by A ij = 1 ( P ij = 0) .

In particular, if P is a stochastic matrix that corresponds to a d -state Markov chain, then starting from any state, we can reach any other state with positive probability in a finite number of steps. The irreducibility property also has an equivalent definition in the matrix form. That is, for any permutation matrix T , TPT -1 cannot be written as an upper triangular block matrix with the following form

In other words, an irreducible matrix does not have a nontrivial absorbing subspace that aligns with the standard basis.

In this work, we require more than the irreducibility property from the transition matrix P π defined in §3.2. In fact, we need the existence of a 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 eigenvalues of P π , denoted by λ , to be bounded away from 1 , which is the leading eigenvalue of the transition matrix. The difference 1 -λ is also referred to as the spectral gap. It is well-known that if all the entries of P π are positive, then P π is irreducible and there is only one leading eigenvalue on the spectral circle with the corresponding eigenvector given by the chain's stationary distribution µ π , and all the other eigenvalues have magnitude strictly less than 1 . However, for our case, the transition matrix P π has zero entries by definition. Fortunately, the nice property on the existence of spectral gap can be generalized to a class called primitive matrix.

Definition C.3 (Primitive Matrix) . A nonnegative and irreducible square matrix P is called primitive if there exists an integer k such that all the entries of P k are positive.

By definition of the primitive matrix, one can immediately see that for any k ′ > k , P k ′ π is a positive matrix. The following is the celebrated Perron-Frobenius theorem that characterizes the spectral structure of the primitive matrices.

Theorem C.4 (Perron-Frobenius Theorem for Primitive Matrices) . Let P be a primitive matrix. Then the following statements hold:

The leading eigenvalue of P is real and positive, and it is the unique eigenvalue with the largest magnitude. In particular, if P is a stochastic matrix, then the leading eigenvalue is 1 .

The leading eigenvector of P is positive and unique up to a scaling factor. In particular, if P is a stochastic matrix, then the leading eigenvector is the stationary distribution of the Markov chain with transition kernel P .

The Perron-Frobenius theorem guarantees the existence of a unique stationary distribution µ π when the transition matrix P π is primitive. In particular, when we further assume that the transition matrix P π has a spectral gap, the chain is sufficiently mixed, meaning that we can thus approximate sum over the entire sequence with an 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):

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

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:

Due to 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 a different sequence length L ′ , the model can still handle the task in the same manner.

By Assumption 3.5, the chain is sufficiently mixed for large L . In the analysis, we actually use X l -M : l = ( x l , x l -1 , . . . , x l -M ) ∼ µ π , where µ π is the stationary distribution over a length-( M +1) window, to approximate the aggregation over X l -M : l for l = 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 is the approximation error. In particular, for sufficiently large L , the CE loss at large l constitutes the majority of the sequential CE loss in (C.5), making the CE loss at small l negligible.

C.5 Standard χ 2 -Divergence and Mutual Information

The χ 2 -divergence (or χ 2 -distance) between two probability distributions P and Q in the same probability space is defined as:

where the summation is taken over all elements x in the sample space where Q ( x ) > 0 . The χ 2 -mutual information between two random variables X and Y with joint distribution P XY and marginal distributions P X and P Y is defined as:

where P X ⊗ P Y is the product of the marginals, meaning ( P X ⊗ P Y )( x, y ) = P X ( x ) P Y ( y ) . For a Markov chain X → Y → Z , the χ 2 -mutual information satisfies the data processing inequality

which follows from the observation that χ 2 -divergence is also an f -divergence.

C.6 More Details on the Generalized Induction Head Mechanism

Recall that we define the Generalized Induction Head (GIH) estimator in (3.3). Specifically, GIH ( x 1: L ; M,D ) is constructed in two steps. First, we find the information-optimal subset S ⋆ of [ M ] by solving (3.2). Second, we build a d -class kernel classifier to predict x L +1 , where the 'data' used by such a classifier are { ψ S ⋆ ( l ) , x l } l ∈ [ M +1 ,L ] . Here { ψ S ⋆ ( l ) , l ∈ [ M +1 , L +1] } are features constructed at each position based on the partial history given S ⋆ . In particular, similar to (C.3), for any subset S of [ M ] , any input token sequence x 1: L , and any position l ∈ [ M +1 , L +1] , we define ψ S ( l ) = ψ S ( l ; x 1: L ) as

In other word, ψ S ( l ) is given by expanding the rank1 tensor spanned by { x l -s } s ∈S into a vector. Here x l -s ∈ X is a vector in R d and we let ( x l -s ) i s denote its i s -th entry. The rationale behind ψ S ( l ) is similar to ϕ introduced in (C.3). We form a long vector containing all the products of the entries of vectors { x l -s } s ∈S . Here we omit the dependency of ψ S on the input sequence x 1: L to simplify the notation. Furthermore, ψ S induces a polynomial kernel such that for any l, m ∈ [ M +1 , L +1] , we have

That is, feature ψ S selects the token position pairs ( l, m ) such that the partial histories induced by S at position l and m are exactly the same.

Based on { ψ S ⋆ ( l ) , x l } l ∈ [ M +1 ,L ] , GIH forms a kernel classifier using the indicator kernel. Specifically, for any j ∈ [ d ] , by (3.3), GIH ( x 1: L ; M,D ) outputs each e j ∈ X with probability

C.7 Further Discussions on the GIH Mechanism

We conclude this section with further discussions on the modified χ 2 -mutual information and lowdegree polynomial kernel for the FFN within the GIH mechanism.

On the Modified χ 2 -Mutual 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 a symmetric case where the stationary distribution µ π over a lengthr n window is uniform

over X r n . One can verify that in this case, the stationary distribution over a window of any other length is uniform as well, and the modified mutual information can be simplified into

where I χ 2 ( S ) is the standard χ 2 mutual information between µ π ( z | Z -S ) and µ π ( z ) , and the second term |S| log d serves as a 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 .

If n = 1 , i.e., each token only has one parent, then S ⋆ = pa . This is because S ⋆ simultaneously maximizes both terms in (C.6), thus reproducing the results in Nichani et al. (2024).

If n is known a priori and restricting the polynomial kernel to S ∈ [ H ] = n = {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-processing 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-parent tokens that are also informative. The rationale is that with a more complex model, e.g., selecting a large S , the model are able to make more accurate predictions for large L but may endure a large estimation error for small L , as the exact matching X l -S = X L +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 accuracy. In this regard, we have the following corollary.

Corollary C.5. We always have |S ⋆ | ≤ n regardless of the choice of D , where S ⋆ = argmax [ H ] ≤ D log r I χ 2 ( S ) for r I χ 2 ( S ) in (C.6)

The reasoning behind this corollary is as follows. Consider any set S with |S| > n , we have I χ 2 ( S ) ≤ I χ 2 ( pa ) as the true parent set is the most informative. Moreover, since | pa | = n < |S| , S suffers from a larger penalty. As a result, we have log r I χ 2 ( S ) < log r I χ 2 ( pa ) when S has more than n elements. In other words, it is without loss of generality to set D ≤ 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 particular, we prove that the generalized induction head mechanism adopts a modified χ 2 -mutual 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 large-scale transformer models.

Our work opens new directions for developing a rigorous understanding of the transformer models. A natural direction would be that if one can find such a mechanism with standard FFN layer using multi-layer perceptron and standard layer normalization in the more practical transformer model. The intuition is 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 MLP 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 occurs 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 to simplify the model based on our choice of the initialization and the structure of the disentangled transformer. We then proceed to present the main proof ideas for the three stages of the gradient flow dynamics, where the training yields the following behaviors:

Stage I: A unique S ⋆ ∈ [ H ] ≤ D stands out such that the associated parameter c S ⋆ dominates those of the other sets. As a result, p ∗ S ( t ) = c 2 S ∗ ( t ) /C D ( t ) approaches to one.

Stage II: For each h ∈ S ⋆ , σ ( w ( h ) ) approaches a one-hot vector e M +1 -h ∈ R M , where w ( h ) contains the parameters of RPE of the h -th head. During this stage, each head concentrates on copying a particular parent.

Stage III: Finally, a grows and reaches O (log L ) . As a result, the trained model approximately implements the GIH mechanism GIH ( x 1: L ; M,D ) .

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 a simpler form.

For the second attention layer, we write the output as y ⊤ = σ ( as ) X where s := u ⊤ L +1 Mask ( U ⊤ 1: L ) ∈ R 1 × L is the row vector of the similarity scores. Recall from (2.5) that the FFN layer with normalization outputs U = ϕ ( V ) / √ C D ∈ R ( L +1) × d e , and we denote the l -th row of U by u l = ϕ ( v l ) / √ C D . For l = M +1 , . . . , L , the l -th entry of s is given by

and the other entries are all -∞ . By the property of the FFN layer in (2.3) and the definition C D = ∑ S∈ [ H ] ≤ D c 2 S , we can rewrite the above attention score as

̸

Note that under Assumption 3.3, by the definition of ∆ w in (3.6), we have a sufficiently large gap w ( h ) -h -w ( h ) -j for all j = h at initialization. Thus, exp( w ( h ) -h ) ≫ exp( w ( h ) -j ) for all j = h , which implies the following approximation:

̸

This further implies that for l = M +1 , . . . , L , we have

which is a binary value indicating whether the query and the key token's history match on the subset S . Combining (D.1) and (D.2), we obtain the following simplified expression for s l :

where we denote p S = c 2 S / ∑ S∈ [ H ] ≤ D c 2 S for S ∈ [ H ] ≤ D .

In summary, when ∆ w is sufficiently large, v ( h ) l approximately copies the token x l -h . As a result, the attention score s l satisfies

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, { c S } S∈ [ H ] ≤ D , and the weights of the first attention layer, { w ( h ) } H h =1 . The analyses of these two stages have similar structures and contain the following essential steps:

Derive the explicit expression of the dynamics of the weights, via direct calculations.

Unveil the key quantities (related to the modified χ 2 -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.

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 c 2 S ( t ) for each S ∈ [ H ] ≤ D . For convenience, we drop the dependence on t in the sequel.

Recall the output of the model is y = ( σ ( a · s ) X ) ⊤ and the cross-entropy loss function is L (Θ) = E π ∼P ,x 1: L [ ℓ (Θ)] , where ℓ (Θ) can be written as ℓ (Θ) = -⟨ x L +1 , log( y + ε 1 ) ⟩ . We ignore the small constant ε in the following proof sketch for simplicity. We also abbreviate the vector of attention probabilities in the second attention layer as σ ∈ R L .

Calculation of the Dynamics of c 2 S . By a direct calculation for the loss ℓ and s l in (D.1),

Here the vector x L +1 /y is obtained by element-wise division and σ l ( a · s ) is the l -th entry of σ ( a · s ) . Then applying the chain rule, we obtain the following dynamics for c 2 along the gradient flow:

S

Note that here the second term f ( t ) is independent of S , and it will be canceled out when we consider the difference of the derivatives, ∂ t log c 2 S -∂ t log c 2 S ′ , for two sets S , S ′ ∈ [ H ] ≤ D . This is why we focus on the time derivative of log c 2 S .

Relate the Dynamics to the Modified χ 2 -MI by Approximations. Now using the approximation in (D.2) for ∏ h ∈S ⟨ v ( h ) l , v ( h ) L +1 ⟩ , expanding ( x L +1 /y ) ⊤ ( x l -y ) coordinate-wise, and noting that σ l ( a · s ) ≈ 1 / ( L -M ) as we have small a in the second attention layer, we arrive at

where y ( k ) denotes the k -th entry of y and X l -S := ( x l -i : i ∈ S ) denotes the history of x l on the set S , similar for X L +1 -S . Note that y ( k ) ≈ ( L -M ) -1 ∑ L l = M +1 1 ( x l = e k ) ≈ µ π ( e k ) , which follows from the mixing assumption of the Markov chain that allows us to replace the average over l = M +1 , . . . , L by the expectation over the stationary distribution. Also for the same reason, we can replace ( x l , X l -S ) , ( x L +1 , X L +1 -S ) with two independent copies from the stationary distribution µ π , i.e.,

See the approximation from g 2 , S to g 3 , S in §E.2. Indeed, the first term in (D.4) becomes the modified χ 2 -MI, r I χ 2 ( S ) , which is defined in Definition 3.1. This gives rise to the following approximation:

Since the value of f ( 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 ∗ = argmax S∈ [ H ] ≤ D r I χ 2 ( S ) that maximizes the modified χ 2 -MI within [ H ] ≤ D .

Convergence of p S ⋆ . Note that p S = c 2 S / ∑ S ′ ∈ [ H ] ≤ D c 2 S ′ quantifies the contribution of the set S to the feature produced by the FFN layer. Thus, it is the relative growth rate of c 2 S that matters. Towards this end, it follows from (D.5) that, for all S ∈ [ H ] ≤ D \{S ⋆ } ,

Here we recall from (3.5) that ∆ r I χ 2 quantifies the minimal gap between the modified χ 2 -MI of S ⋆ and any other set in [ H ] ≤ D . The lower bound given by (D.6) ensures that for all S ̸ = S ⋆ , the ratio c 2 S ⋆ /c 2 S grows exponentially fast, which further implies that p S ⋆ approaches one exponentially fast. This concludes the first stage of the training dynamics.

D.2.2 Training the First Attention Layer: Convergence of σ ( w ( h ) ) to One-Hot Vector

As we proceed to the second stage after p S ⋆ ≈ 1 , it suffices to show how σ ( w ( h ) ) converges to a one-hot vector e M +1 -h for h ∈ S ⋆ in order to show that the model converges to the GIH mechanism. Recall that we denote X = ( x 1 , . . . , x L ) ∈ R L × d . For notational convenience, we denote σ ( h ) := σ ( w ( h ) ) and let X ( l -M ):( l -1) ∈ R M × d denote the submatrix of X with rows l -M,...,l -1 for any l . Following our convention, we let σ ( h ) -i denote the ( M +1 -i ) -th entry of σ ( h ) and similarly for w ( h ) -i .

Calculation of the Dynamics of w ( h ) . The main idea for analyzing { w ( h ) } H h =1 is the same as that in the previous stage: It suffices to analyze the difference between the growth rates of different coordinates of w ( h ) for h ∈ S ⋆ . In particular, we care about how quickly w ( h ) -h grows compared to other coordinates if w ( h ) -h is initialized to be larger than the remaining coordinates:

Now, we invoke the result obtained in the previous stage that p S ⋆ ≈ 1 , which gives us s l ≈ ∏ h ∈S ⋆ ⟨ v ( h ) l , v ( h ) L +1 ⟩ . Consequently, for any h ∈ S ⋆ , we have

where the equality follows from the fact that w ( h ) -i only affects ( v ( h ) l , v ( h ) L +1 ) and differentiating through the softmax function. Here we define b l := X ( l -M ):( l -1) v ( h ) L +1 + X ( L +1 -M ): L v ( h ) l to simplify the notation. Combining (D.7) and (D.8), we obtain

̸

where we introduce the following notation

̸

Adetailed deviation of (D.9) can be found in (E.16). Notice that σ ( h ) -h -σ ( h ) -i is positive at initialization. Now suppose σ ( h ) -h -σ ( h ) -i > 0 holds at current time t . Then, lower bounding ∂ t w ( h ) -h -∂ t w ( h ) -i boils down to lower bounding g ⊤ h ( e M +1 -h -e M +1 -i ) for i = h . Furthermore, if we can show that ∂ t w ( h ) -h -∂ t w ( h ) -i is lower bounded by some positive value, the gap σ ( h ) -h -σ ( h ) -i will further increase. Since ∑ M i =1 σ ( h ) -i ≡ 1 , this will create a reinforcing loop that makes σ ( h ) -h monotonically increase.

̸

̸

Relate the Dynamics to the Modified χ 2 -MI by Approximations. We demonstrate next that g ⊤ h ( e M +1 -h -e M +1 -i ) for i = h admits a lower bound depending on the information gap ∆ r I χ 2 . Specifically, using the same strategy for (D.3), we have by definition that

where for b l we have by the same approximation v ( h ) l ≈ x l -h and v ( h ) L +1 ≈ x L +1 -h as in (D.2) that

̸

Now we consider the case i = h and i = h separately:

(i) ( i = h ) For g ⊤ h e M +1 -h , we simply set i = h in (D.11), and the indicator 1 ( x L +1 -h = x l -h ) will exactly compensate for the exclusion of h in the indicator function of (D.10). Drawing an analogy to how we go from (D.3) to (D.5), we obtain

(ii) ( i = h ) For g ⊤ h e M +1 -i with i = h in (D.11), we apply the same reasoning as in the previous case. Additionally, by using the Cauchy-Schwarz inequality, the following inequality holds up to a small error (see Lemma F.7 for a detailed derivation):

̸

̸

Plugging this back into the dynamics in (D.9), we conclude that for all i = h ,

Convergence of σ ( w ( h ) ) . Combining the arguments in the previous two steps, we can now say that σ ( h ) -h will monotonically increase. It remains to show that σ ( h ) -h converges to one. Note that log( σ ( h ) -h /σ ( h ) -i ) = w ( h ) -h -w ( h ) -i by the definition of the softmax function. Therefore,

where σ ( h ) -h (0) is the initial value of σ ( h ) -h at time t = 0 . One can now rearrange the term and pick the ratio σ ( h ) -i /σ ( h ) -h as the variable to track in the dynamics. A refined analysis in the convergence analysis in §E.3 shows that σ ( h ) converges to a one-hot vector with σ ( h ) -h going to one. In particular, the convergence rate is determined by the information gap ∆ r I χ 2 according to the above formula.

D.3 Analysis for the Training of the Second Attention Layer

In the last stage, we turn to the training of a given that all σ ( h ) 's for h ∈ S ⋆ are approximately one-hot vectors. The following approximation of the dynamics of a ( t ) is performed in the region a ≤ O (log L ) , where the signal term in the dynamics dominates the approximation error.

Calculation of the Dynamics of a . After Stages I and II, the output is approximated as y ( k ) ≈ y ⋆ ( k ) := ∑ L l =1 σ ⋆ l 1 ( x l = e k ) for each k ∈ [ d ] . Here the weighting coefficients σ ⋆ 1 , . . . , σ ⋆ L satisfy

Note that for each l ∈ [ L ] , σ ⋆ l indicates the importance assigned to the l -th token based on the corresponding history of x l over the information set S ⋆ . In the population counterpart, when the chain has sufficiently mixed, for given X L +1 -S ⋆ , we can roughly view each ( x l , X l -S ⋆ ) as being sampled from a reweighed version of the stationary distribution :

Following the same argument as those in the previous stages, replacing the sum over l with the expectation over the stationary distribution, we arrive at

See detailed derivations of the above approximation in §E.4. Comparing the above expression with (D.4) in Stage I, one can see that here ( x, X -S ⋆ ) and ( z, Z -S ⋆ ) are no longer independent because now the model has learned to perform a information-theoretic 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 χ 2 -mutual information. In fact, the underlying joint distribution q π is given by q π = µ π ( x, X -S ⋆ ) · r µ π ( z, Z -S ⋆ | X -S ⋆ ) .

Divergence of a . As the dynamics of a has no closed-form expression due to the nonlinearity in the reweighed distribution r µ π , 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 superexponential growth until it reaches a critical 'elbow' value. After that, when a becomes even larger, it grows logarithmically until it reaches Ω(log L ) .

E Analysis of the Training Dyanamics

Masking the Simplified Model. Recall that we apply a mask to the first M position in the simplified model. Therefore, we only allow index l to run from M +1 to 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 :

where Ω only hides a universal constant that does not depend on the model parameters. The conditions in (E.1) will facilitate our analysis for Stage I and Stage II. For the last stage, we require

where

and Ω only hides universal constants that do not depend on the model parameters. Here, µ π ( x, X -S ⋆ ) denotes the stationary distribution of the Markov chain over token x and its parents X -S ⋆ , with S ⋆ being the information set defined in (3.2).

E.2 Analysis for Stage I

In this section, we analyze the dynamics of the parameters { c 2 S } S∈ [ H ] ≤ D in the first stage of training. We will show that there is a unique S ∗ ∈ [ H ] ≤ D such that c 2 S ⋆ dominates all the other c 2 S 's 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 a high level, the strategy is to analyze ∂ t log c 2 S ⋆ -∂ t log c 2 S for all S ̸ = S ⋆ via the following steps:

Dynamics Calculation. First, we calculate the dynamics of log c 2 S for each fixed S . By selecting sufficiently small values for a and ε , and leveraging the mixing properties of the Markov chain with large L , the dynamics of log c 2 S is approximately governed by the modified mutual information r I χ 2 ( S ) .

Lower Bound for The Growth Rate. Consequently, we are able to lower bound the difference between the growth rates, ∂ t log c 2 S ⋆ -∂ t log c 2 S , in terms of ∆ r I χ 2 , the gap between the modified mutual information of S ⋆ and the second-best set.

Convergence. Finally, we derive the convergence using the above lower bound.

Before presenting the proof, we first remind the readers of a few definitions and notations. Recall that our simplified model is given by

Also recall that C D ( t ) = ∑ S∈ [ H ] ≤ D c 2 S ( t ) and p S ( t ) = c 2 S ( t ) /C D ( t ) for each S ∈ [ H ] ≤ D . The loss function can be rewritten as

Here the expectation E is taken over both the sequence ( x 1 , . . . , x L +1 ) and the Markov kernel π ∼ P . We abbreviate σ ≡ σ ( as ) for convenience and denote by σ l the l -th element of σ . By direct calculation, we have

Then applying the chain rule, we have

In addition,

Now, we are ready to present the proof of Theorem 3.6 for the first stage of training. We remind readers that here only { c S } S [ H ] ≤ D are trained, and we omit the dependence on t for convenience.

Proof of Theorem 3.6: Stage I. As discussed in the proof strategy above, we first derive the dynamics of log c 2 S for each fixed S ∈ [ H ] ≤ D . Then we compare the growth rate of c 2 S ⋆ with any other c 2 S .

Calculation of The Dynamics of log c 2 S . We fix a S ∈ [ H ] ≤ D and apply the chain rule ∂ℓ/∂c S = ∑ L l = M +1 ∂ℓ/∂s l · ∂s l /∂c S and the gradient flow formula that ∂ t c 2 S = -2 c S · ∂ L /∂c S . We have

In the following, we consider a fixed π for error analysis and take expectation over π again when plugging in everything back into the dynamics. To simplify the expression of ∂ t c 2 S , we define quantities g 0 , S and f as

Note that here f does not depend on S . Based on the above definitions, we can rewrite (E.4) as

Using this, it can be shown that C D ( t ) does not change during the training, as described in the following lemma.

Lemma E.1. The quantity C D ( t ) = ∑ S∈ [ H ] ≤ D c 2 S ( t ) is preserved along the gradient flow over { c S } S∈ [ H ] ≤ D , i.e., ∂ t C D ( t ) ≡ 0 .

This lemma will be useful in the following analysis, and we defer its proof to §E.2.1. Next, we proceed to further simplify the dynamics in (E.5) by approximating g 0 , S .

Simplification of ∂ t log c 2 S . To approximiate g 0 , S , we introduce the following quantities:

Here s y ( k ) is the k -th entry of s y . We remark that each of g 1 , S , g 2 , S , g 3 , S is a function of π and t , but we omit the dependence for brevity.

From g 0 , S to g 1 , S , we replace attention probability σ l ( as ) by the uniform average with factor 1 /L , which yields s y . From g 1 , S to g 2 , S , we replace the empirical distribution s y with the stationary distribution µ π and drop the small constant ε . Finally, from g 2 , S to g 3 , S , we replace the average over the sequence by the expectation over the stationary distribution µ π of the underlying Markov chain. We will show that the approximation error in each step is small, given that a and ε are sufficiently small and the Markov chain mixes well for a large L .

For the approximation of g 0 , S by g 1 , S , note that when a is small, the attention probability σ l ( as ) ≈ 1 / ( L -M ) for all l ∈ [ L ] . More specifically, it follows from Lemma F.3 that

For the approximation of g 1 , S by g 2 , S , we leverage the approximation s y ( k ) ≈ µ π ( e k ) due to the mixing of the Markov chain for large L . The result in Lemma F.4 implies that

where µ 0 ( · ) is the initial distribution over the first r n tokens. Here we abuse the notation of µ π in D χ 2 ( µ 0 ∥ µ π ) to denote the stationary distribution over the last r n tokens. Since µ π min ≥ γ by Assumption 3.5, we have

Therefore, we can further simplify the above bound as

Finally, the approximation of g 2 , S by g 3 , S follows from the mixing property of the Markov chain. In particular, it follows from Lemma F.5 that

Combining the above results, and by the assumption that a = a (0) = O (1 /L 3 / 2 ) and ε = 1 / √ L , we obtain the following approximation error:

where we note that |S| ≤ D for any S ∈ [ H ] ≤ D and the last inequality holds by also noting our condition on L in (E.1) that L ≥ Ω((1 -λ ) -1 γ -D ) . As a result, the dynamics of c 2 S in (E.5) can be approximated as follows:

where O ( · ) hides universal constants that do not depend on the model parameters. Here and in the sequel, we let E denote an error term that is of the order O ( a/ √ C 2 D L (1 -λ ) γ r n +2 ) where the specific constant hidden in O ( · ) may change from line to line, but does not depend on the model parameters. In fact, we can show C D remains constant by Lemma E.1 and a is not updated during this stage. Thus, the error term |E| is of scale O ( aL -1 / 2 ) .

Lower Bound for The Difference ∂ t log c 2 S ⋆ -∂ t log c 2 S . The reason for approximating g 0 , S by g 3 , S in the previous step is that the latter is more interpretable, in the sense that we can relate it to the modified χ 2 mutual information r I χ 2 ( S ) . Recall that for each S ∈ [ H ] ≤ D , the modified χ 2 -mutual information is

Note that f in (E.7) is independent of S , and will be canceled when computing ∂ t log c 2 S ⋆ -∂ t log c 2 S :

Thus, it suffices to consider E π ∼P [ g 3 , S ⋆ -g 3 , S ] . It follows from Lemma F.6 that for each S ∈ [ H ] ≤ D , E π ∼P [ g 3 , S ] satisfies

This yields a lower bound for E π ∼P [ g 3 , S ⋆ ] and an upper bound for E π ∼P [ g 3 , S ] for each S ̸ = S ⋆ , i.e.,

Consequently,

where the second inequality follows from the definition ∆ r I χ 2 = min S ∈ [ H ] ≤ D \{S ⋆ } r I χ 2 ( S ⋆ ) -r I χ 2 ( S ) . Moreover, since each ( σ ( h ) -h ) 2 ∈ (0 , 1) , we have ∏ h ∈S ( σ ( h ) -h ) 2 ≥ ∏ H h =1 ( σ ( h ) -h ) 2 for any S ∈ [ H ] ≤ D . Appling this to the above inequality, we obtain

Exponential Growth of c 2 S ⋆ . We proceed to show that the first term in (E.8) dominates the error term E and thus leads to the exponential growth of c 2 S ⋆ .

̸

Note that by Assumption 3.3, w ( h ) -h ≥ w ( h ) -j +∆ w for all j = h and h ∈ [ H ] , where the quantity ∆ w satisfies

Recall that we are not updating the RPE parameters during this stage, so σ ( h ) is fixed for all h ∈ [ H ] . So the gap condition (E.9) holds throughout Stage I. This conditions ensures that w ( h ) -h ≫ w ( h ) -j , so ∏ h ∈ [ H ] ( σ ( h ) -h ) 2 is sufficiently large. More precisely, given that head h is more focused on the ( -h ) -th position by having a gap ∆ w in the initialization, we can further show by definition of the softmax function that

Plugging (E.9) into (E.10), we have by additionally noting that r I χ 2 ( S ⋆ ) ≥ ∆ r I χ 2 > 0 that

which implies that

Moreover, when L is sufficiently large such that L ≥ Ω((∆ r I 2 χ 2 (1 -λ ) γ r n +2 ) -1 ) , E in (E.8) satisfy |E| ≤ 13 a ∆ r I χ 2 / 6 C D , where Ω hides a universal constant that does not depend on the model parameters. Therefore, combining (E.8) and (E.11), we conclude that

This implies that c 2 S ⋆ grows exponentially fast and becomes dominant.

Convergence of p S ⋆ . In this part, we treat all the model parameters as a function of time t . For simplicity, we omit the dependence on t when it is clear from the context. It remains to derive the convergence of p S ⋆ = c 2 S ⋆ /C D . Expanding C D = ∑ S∈ [ H ] ≤ D c 2 S , we can directly calculate the derivative of p S ⋆ as follows:

where in the last equality we use the fact that ∂ t log c 2 S = ( ∂ t c 2 S ) /c 2 S . Applying (E.12) to each S ̸ = S ⋆ , we further have

This implies that p S ⋆ = c 2 S ⋆ /C D monotonically increases, and thus c 2 S ⋆ ( t ) ≥ c 2 S ⋆ (0) for any t ≥ 0 because C D is constant by Lemma E.1 and c 2 S ⋆ (0) is the initial value for c 2 S ⋆ at time t = 0 . Therefore, we can further replace c 2 S ⋆ by its initial value in the above inequality, which yields

We remark that the above upper bound is independent of t . Finally, applying the Grönwall's inequality to log(1 -p S ⋆ ) , we obtain

With training time t 1 ≥ (2 C D (0) log L ) / ( a · p S ⋆ (0)∆ r I χ 2 ) , we can guarantee that

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.1.

Proof of Lemma E.1. By (E.5), we have

Moreover, by the definition of g 0 , S and f , it holds that ∑ S∈ [ H ] ≤ D p S g 0 , S = f . Then,

Thus, the quantity C D is preserved under the dynamics.

E.3 Analysis for Stage II

In this section, we provide the analysis of the dynamics of σ ( h ) ≡ σ ( w ( h ) ) for head h ∈ S ⋆ . For head h / ∈ S ⋆ , the results from Stage I imply that p S → 0 for any S ̸ = S ⋆ . Consequently, any head h / ∈ S ⋆ will be ignored when producing the output features of FFN. Conversely, for h ∈ S ⋆ , we establish the dominance of w ( h ) -h over w ( h ) -i for all i = h , yielding σ ( h ) -h → 1 as t →∞ . 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, ∂ t w ( h ) -h -∂ t w ( h ) -i for all i = h , via the following steps:

Dynamics Calculation. We initiate the analysis by deriving the dynamics of w ( h ) -i for any fixed i and h .

Dynamics Approximation Then we approximate the dynamics by identifying the dominant term controlled by the modified χ 2 mutual information r I χ 2 ( S ⋆ ) .

̸

Lower Bound for The Growth Rate By comparing the corresponding modified χ 2 mutual information, we establish a lower bound on ∂ t w ( h ) -h -∂ t w ( h ) -i for all i = h .

Convergence. Finally, we derive the convergence rate of σ ( h ) -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 ) } H h =1 . For the h -th head of the first attention layer, the positional embedding vector w ( h ) induces the attention probability over a window of size M , i.e.,

Further recall the attention scores for the second attention layer, as , where s = u ⊤ L +1 U ⊤ M +1: L . Then for each l ∈ [ L ] , the l -th coordinate of s is given by

Here p S is defined as in the analysis of Stage 1, and X ( l -M ):( l -1) ∈ R M × d is the submatrix of X with rows l -M,...,l -1 .

By direct calculation, we have

Then by chain rule,

Moreover, we can view each s l as a function of { v ( h ) 1 , . . . , v ( h ) L +1 } h ∈ [ H ] . Differentiating s l with respect to v ( h ) l and v ( h ) L +1 , we have

In the summation, we only add those S 's in [ H ] ≤ D containing h . Also, recall from (E.3) that

Now we are ready to proceed with the analysis for Stage II.

Proof of Theorem 3.6: Stage II. We start by calculating the explicit expression of the dynamics of ∂ t w ( h ) -i , and then derive approximation of the dynamics, which allows us to further show the convergence of σ ( h ) .

Calculation of The Dynamics of ∂ t w ( h ) . First fix an h ∈ [ H ] . To simplify the notation, for each l ∈ [ L ] we define

Note that w ( h ) is the parameters of the h -th head and only enters each v ( h ) l , l = 1 , . . . , L +1 . Recall that s l = ∑ S∈ [ H ] ≤ D p S ∏ h ∈S ⟨ v ( h ) l , v ( h ) L +1 ⟩ , and the RPE weight w ( h ) for attention head h only influences its outputs v ( h ) l and v ( h ) L +1 in the sum. It thus follows from the chain rule that for each i ∈ [ M ] , we have

where we remind readers that e i ∈ R M × 1 is the i -th standard basis vector.

Furthermore, along the gradient flow ∂ t w ( h ) -i = -∂ L /∂w ( h ) -i , it follows from (E.13) that

where we plug in the expression of ∂s l /∂w ( h ) -i above in the last equality. Here the vector g h, 0 is defined as

where σ l is the softmax probability for the l -th token in the second attention layer. Comparing ∂ t w ( h ) -i and ∂ t w ( h ) -h , we have

Using the fact that ∑ M j =1 σ ( h ) -j = 1 , we can rewrite

where in the first identity, we add and then subtract term σ ( h ) -i e M +1 -h . Combining (E.15) and (E.16) yields for each i ∈ [ M ] that

Simplification of ∂ t w ( h ) -i . We proceed by deriving approximations to the vector g h, 0 , which will help us identify the dominant term in the dynamics ∂ t w ( h ) -h -∂ t w ( h ) -i . Specifically, we define

, where Z = [ z -M , . . . , z -1 ] ⊤ ∈ R M × d is an independent copy of X = [ x -M , . . . , x -1 ] ⊤ ∈ R M × d , and v ( h ) ( X ) := M ∑ i =1 σ ( h ) -i x -i , v ( h ) ( Z ) := M ∑ i =1 σ ( h ) -i z -i , b ( X,Z ) := Z ( v ( h ) ( X )) + X ( v ( h ) ( Z )) , s y := 1 L -M L ∑ l = M +1 x l .

The strategy of gradually approximating g h, 0 by g h, 1 , g h, 2 , g h, 3 and g h, 4 is similar to the analysis in Stage I. To see the intuition, from g h, 0 to g h, 1 , we use the fact that p S ⋆ ≈ 1 and p S ≈ 0 for any other S , which is a result of Stage 1. From g h, 1 to g h, 2 , we replace y by the empirical mean s y , thanks to the fact that σ l ( a ) ≈ 1 /L when a is small. Then, from g h, 2 to g h, 3 , we replace the empirical distribution s y with the stationary distribution of the Markov chain. These two steps also appear in the analysis of Stage 1. Finally, to go from g h, 3 to g h, 4 , we leverage the rapid mixing of the Markov chain.

̸

Note that the common structures in (E.15) are g ⊤ h, 0 ( e M +1 -h -e M +1 -i ) for i = h . Hence, we only need to understand the approximation error in each step for g ⊤ h, 0 ( e M +1 -h -e M +1 -i ) . Recall that we are focusing on h ∈ S ⋆ in this stage.

From g h, 0 to g h, 1 , we remove the terms in the summation that are weighted down by p S for any S ̸ = S ⋆ due to the rapid dominance of p S ⋆ from Stage I. Recall that p S ∗ converges to one at an exponential rate while all other p S 's converge to zero. For simplicity, let us define

By the triangular inequality, we have

where in the last line we use the claim that | ρ ( S ) | ≤ 8 for all S . To see this point, note that by definition of b l in (E.14), we have

since v ( h ) l and x l have norm at most 1 . Then, by Lemma F.2 where we plug in the upper bound 4 for the function f ( · ) in the lemma, we conclude that ρ ( S ) ≤ 8 , ∀S ∈ [ H ] ≤ D \ {S ⋆ } . Define ∆ 1 := 1 -p S ⋆ ( t 1 ) , and ∆ 1 ≤ 1 /L by the results from Stage I. Thus, we obtain

For the approximation of g h, 1 by g h, 2 , we use the fact that σ l ( as ) ≈ 1 /L when a is sufficiently small. Specifically, we also take the absolute bound for f ( · ) as 4 in Lemma F.3 and obtain

For the approximation of g h, 2 by g h, 3 , we use the fact that s y ( k ) ≈ µ π ( e k ) for large L . More precisely, it follows from Lemma F.4 with the upper bound 4 for f ( · ) in the lemma that

Finally, to go from g h, 3 to g h, 4 , we leverage the rapid mixing of the Markov chain. Intuitively, when l and L +1 are far apart, x l and its parents in S ⋆ are independent of x L +1 and its parents in S ⋆ . This observation yields the approximation of g ⊤ h, 3 ( e M +1 -h -e M +1 -i ) by g ⊤ h, 4 ( e M +1 -h -e M +1 -i ) . To simplify the notation, define two scalars

Using the notation above, we have

Recall that

We apply the triangular inequality to obtain that

Each term on the right-hand side can be bounded by Lemma F.5, where in the lemma we take ( σ ( h ) ) h ′ ∈S ⋆ ∈ R M ×|S ⋆ | and (( σ ( h ′ ) ) h ′ ∈S ⋆ \{ h } , e h ) ∈ R M ×|S ⋆ | as the two lists of vectors on the M -dimensional probability simplex for r σ and σ respectively. Consequently, we have

Combining the above results and setting ε = 1 / √ L , a = a (0) ≤ O (1 /L 3 / 2 ) and together with the conditions in (E.1), we have

where O ( · ) hides universal constants independent of the parameters of the model. We remark that while the left hand side is a function of t , the upper bound is independent of t . Then, we can rewrite (E.17) as

̸

̸

Lower Bound for The Difference ∂ t w ( h ) -h -∂ t w ( h ) -i . To show ∂ t w ( h ) -h -∂ t w ( h ) -i > 0 , we first derive the lower bound of E π ∼P [ g ⊤ h, 4 ( e M +1 -h -e M +1 -i )] for any i = h . Since ( x, X ) and ( z, Z ) are independent and identically distributed, by the definition of b ( X,Z ) ,

where we introduce the following quantities for convenience:

The quantities τ h, 1 and τ h, 2 can be further approximated. Specifically, by applying Lemma F.6 to τ h, 1 , , where in the lemma we take ( σ ( h ) ) h ′ ∈S ⋆ ∈ R M ×|S ⋆ | and (( σ ( h ′ ) ) h ′ ∈S ⋆ \{ h } , e h ) ∈ R M ×|S ⋆ | as the two lists of vectors on the M -dimensional probability simplex for σ and r σ respectively, and we obtain

Drawing on the analagous reasoning as in the proof of Lemma F.6, we can approximate τ h, 2 as follows:

where

To establish the lower bound for τ h, 1 -τ h, 2 , let us begin by establishing an upper bound for ψ , which is approximately equal to τ h, 2 . We invoke Lemma F.7 with S = S ⋆ and S ′ = S ⋆ \{ h } ∪ { i } in the lemma to obtain

̸

Leveraging this for (E.19) and (E.20),

where in the second line we multiply an additional σ ( h ) -h to the product as σ ( h ) -h ∈ [0 , 1] .

Next, we provide a lemma showing that ∂ t σ ( h ) -h is growing for all time t ≥ t 1 , where t 1 is the starting time of the second stage.

̸

Lemma E.2 (Reinforced Growth of σ ( h ) -h ) . For all h ∈ S ⋆ , we have for all i = h at any t ≥ t 1 :

Proof. See §E.3.1 for the proof.

In the proof of Lemma E.2, we will use the following useful proposition.

̸

Proposition E.3. Suppose ∏ h ∈S ⋆ ( σ ( h ) -h ) 2 ≥ 1 / (1 + ( M -1) exp( -∆ w )) 2 |S ⋆ | with ∆ w satisfying (3.6) , and σ ( h ) -h > σ ( h ) -i for any i = h, h ∈ S ⋆ at a given time t . Suppose Assumption 3.5 holds and L satisfies (E.1) . It holds that

̸

̸

̸

Lemma E.2 implies that during Stage II, for all i = h and h ∈ S ⋆ , we have w ( h ) -h > w ( h ) -i and σ ( h ) -h > σ ( h ) -i for all t ≥ t 1 . In addition, as σ ( h ) -h is growing, all the conditions in Proposition E.3 are satisfied for any t ≥ t 1 , and hence all the conclusions in (E.23).

Convergence of σ ( h ) . Finally, we characterize the convergence rate of σ ( h ) . For the convergence analysis, we adhere to the convention used in the previous stage, treating all model parameters as functions of the training time t , where t = t 1 marks the start of the second stage. With a slight abuse of notation, we denote by σ ( h ) -i ( t ) the value of σ -i ( w ( h ) ( t )) at time t , where w ( h ) ( t ) is the input to the softmax function, and σ -i ( · ) refers to the ( M +1 -i ) -th element of the softmax probability. For simplicity, we sometimes omit the time index t when the context makes it clear.

Note that ∂ t σ ( h ) -h > 0 for all h ∈ S ⋆ . Hence by the definition of the softmax operation, we have

̸

where the first inequality follows from the monotone growth of σ ( h ) h , and the second line follows from the fact that the first attention layer is untouched during the first stage. Note that here in (E.24), σ ( h ) and w ( h ) are functions of t . Now, putting together (E.23) and (E.24), and also noting that σ ( h ) -h > σ ( h ) -i for all i = h and h ∈ S ⋆ , it follows that

Rearranging the terms, and using the fact that w ( h ) -h ( t 1 ) -w ( h ) -i ( t 1 ) ≥ ∆ w by Assumption 3.3, we get

This yields a lower bound for σ ( h ) -h ( t ) as follows:

̸

where we define σ min (0) := min h ∈S ⋆ σ ( h ) -h (0) . Consequently, we have

Now, we consider large t such that

Then, we can apply the inequality (1 -x ) n ≥ 1 -nx for x ∈ [0 , 1 /n ] and n ≥ 1 to obtain

Therefore, with training time t 2 = 4 L |S ⋆ | · ( M -1) /a ∆ r I χ 2 · σ min (0) + t 1 , we can ensure that

This completes the proof for Stage II.

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 ∏ h ∈S ⋆ ( σ ( h ) -h ) 2 ≥ 1 / (1 + ( M -1) exp( -∆ w )) 2 |S ⋆ | with ∆ w in (3.6) implies that

Combining (E.21) and (E.25) yields

̸

for any i = h . Applying (E.26) to (E.18), since each σ ( h ) -i > 0 and σ ( h ) -h > σ ( h ) -i at time t for all i = h, h ∈ S ⋆ , it holds that

̸

̸

Then since we assume a sufficiently large L ≥ Ω((∆ r I 2 χ 2 (1 -λ ) γ r n +2 ) -1 ) , it holds that |E| ≤ ∆ r I χ 2 / 6 , we further have

̸

̸

As ∂ t log σ ( h ) -h -∂ t log σ ( h ) -i = ∂ t w ( h ) -h -∂ t w ( h ) -i > 0 by property of the softmax function, and ∑ M i =1 ∂ t σ ( h ) -i = 0 , we have ∂ t σ ( h ) -h > 0 for all h ∈ S ⋆ . This completes the proof of Proposition E.3.

̸

Proof of Lemma E.2. We give a proof to Lemma E.2 by contradiction. Note that at the beginning of the second stage t = t 1 , we have all the conditions for Proposition E.3 satisfied by the initialization conditions in Assumption 3.3. Then, by (E.23) in Proposition E.3, we have ∂ t log σ ( h ) -h -∂ t log σ ( h ) -i > 0 and ∂ t σ ( h ) -h > 0 for all i = h and h ∈ S ⋆ at t = t 1 .

̸

Next, assume that τ > t 1 is the smallest time such that at least ∂ t σ ( h ) -h ≤ 0 or ∂ t log σ ( h ) -h -∂ t log σ ( h ) -i ≤ 0 for some i = h and h ∈ S ⋆ . By definition of τ , we have (E.22) holds for any moment t ∈ [ t 1 , τ ) . As σ ( h ) -h and the gap σ ( h ) -h -σ ( h ) -i are monotonically increasing, we have by the initialization condition and the boundedness of the gradient that at time τ :

̸

̸

Hence, by Proposition E.3, we have ∂ t log σ ( h ) -h -∂ t log σ ( h ) -i > 0 and ∂ t σ ( h ) -h > 0 for all i = h and h ∈ S ⋆ at time τ , which contradicts the definition of τ . This completes the proof of Lemma E.2.

E.4 Analysis for Stage III

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 a < O (log L ) , 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.

Proof Strategy. We analyze the dynamics of a via the following steps:

Dynamics Calculation. First, we derive the explicit expression for the dynamics of a .

Dynamics Approximation. We approximate the dynamics by exploiting the mixing properties of the Markov chain and the convergence of the weights from Stage I and II.

Lower and Upper Bound for The Growth Rate. Finally, we establish the upper and lower bounds for the growth rate of a when a is either sufficiently small or large.

For a set S ⊆ [ M ] , we denote X l -S := ( x l -s : s ∈ S ) . If l = 0 , we will ignore l in the subscript and simply use X -S . In this section, we abbreviate p S ( t 1 ) after the first stage's training as p S , and σ ( h ) -i ( t 2 ) after the second stage's training as σ ( h ) -i .

Proof of Theorem 3.6: Stage III. We start with the explicit expression of the dynamics of a .

Calculation of The Dynamics of a . First by the chain rule,

where in the last equality we remind readers of the same procedure as we have used in the derivation of (E.3) in Stage I. Then, taking expectation with respect to X and π and expanding s l = a ∑ S∈ [ H ] ≤ D p S ∏ h ∈S ⟨ v ( h ) l , v ( h ) L +1 ⟩ , we have

We remind readers the shorthand σ ≡ σ ( as ) . We denote the above quantity by f 0 .

Approximation of ∂ t a . Similar to the analysis for the previous two stages, we develop a sequence of approximation steps that transforms ∂ t a into a tractable quantity. We aim to decouple x L +1 and x l , approximate s l by a 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 [ H ] ≤ D \{S ⋆ } where S ⋆ is the optimal set that maximizes the modified mutual information defined in (3.1). This is because c S ⋆ dominates by the analysis of Stage I. Specifically, we define

To bound | f 0 -f 1 | , note that for any S ∈ [ H ] ≤ D , since each v ( h ) l has norm at most 1 , we can invoke Lemma F.2 with C = 1 and obtain

It follows that

In summary, the difference between f 0 and f 1 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 in the second attention layer. Let us define s ⋆ l = ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) as the ideal attention score for the second attention layer. We invoke Lemma F.1 to have for all l ∈ [ L ] ,

Corresponding to { s ⋆ l } L l = M +1 , we define

In the vector form, we have y ⋆ = ∑ L l = M +1 σ ⋆ l x l . Leveraging the above approximations, we define an approximation of f 1 as

Applying Lemma F.9, it holds that

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 ) is also random due to the randomness in σ ⋆ l , and as L is sufficiently large, we want to replace y ⋆ ( k ) with its population counterpart. Let z ∈ X and Z = ( z -M , . . . , z -1 ) ∈ X M be two random variables and we define similarly for x ∈ X and X = ( x -M , . . . , x -1 ) ∈ X M . To this end, we define a reweighed distribution

where µ π is the stationary distribution of the Markov chain over a window of size M +1 . This can be viewed as a reweighting of the stationary distribution over ( z, Z ) by an exponential term that depends on the sequence X -S ⋆ . We use r µ π ( z = e k | X L +1 -S ⋆ ) to replace y ⋆ ( k ) and define f 3 as

One can immediately draw a connection to Lemma F.4 as 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 f 2 and f 3 as

where µ 0 ( · ) is the initial distribution for the first r n tokens in the Markov chain. Here and in the sequel, we simply use D χ 2 ( µ 0 ∥ µ π ) to denote D χ 2 ( µ 0 ( X 1: r n = · ) ∥ µ π ( X 1: r n = · )) when it is clear from the context. In the last inequality, we use the fact that D χ 2 ( µ 0 ∥ µ π ) ≤ γ -r n by (E.6) and the condition ε = L -1 / 2 .

Note that in the expression of f 3 , each σ ⋆ l still implicitly depends on the actual value of the sequence X . Since L is large and the Markov chain is well-mixed, we can approximate

∑ L l = M +1 σ ⋆ l 1 (( x l , X l -S ⋆ ) = ( · , · )) by r µ π ( · , · | X L +1 -S ⋆ ) . This gives rise to the following approximation of f 3 :

Applying Lemma F.11 yields

where we use the fact that D χ 2 ( µ 0 ∥ µ π ) ≤ γ -r n by (E.6).

Let ( z, Z ) ∼ r µ π ( · | X L +1 -S ⋆ ) . Since L is large, the distribution of ( x L +1 , X L +1 -S ⋆ ) is close to the stationary distribution µ π . Thus, we introduce the following approximation of f 4 :

Note that

and so is | r µ π ( Z -S ⋆ = X -S ⋆ | X -S ⋆ ) | ≤ 1 . The difference between f 4 and f 5 is thus bounded by 2 ∥ p π ( x L +1 , X L +1 -S ⋆ = · , · ) -µ π ( x L +1 , X L +1 -S ⋆ == · ) ∥ TV and by the results in (F.29) of Lemma F.16:

where we use D χ 2 ( µ 0 ∥ µ π ) ≤ γ -r n and the condition on L in (E.2).

Collecting all the above approximation steps, we obtain (where we use ≲ to hide absolute constants)

where the last line holds by moting that with sufficiently large t 1 and t 2 we have ∆ 1 +∆ 2 ≤ L -1 , and ε = L -1 / 2 . Here, express the error in terms of the trainable parameter a and define

In particular, we have for a = O (log L ) that

In a nutshell, we conclude that when the weight a satisfies a < O (log L ) , the dynamics of a can be approximated by

The following proposition helps us reformulate f 5 in a form that facilitates the analysis of the dynamics of a .

Proposition E.4. The term f 5 can be reformulated as

where r π ( X -S ⋆ ; a ) = (1 + µ π ( X -S ⋆ ) · ( e a -1)) -1 is the inverse of the normalization factor of r µ π in (E.27) and

Proof. See §E.4.1 for the proof.

Inspired by this form, we define an alternative function r J ( · ; r, π ) as

where we replace r π ( X -S ⋆ ; a ) by a parameter r ∈ [0 , 1] . As exactly calculating the inverse normalization factor r π ( X -S ⋆ ; a ) is intractable, we instead seek to find an upper and lower bound for r π ( X -S ⋆ ; a ) and plug them into r J ( · ; r, π ) to bound f 5 Suppose that r π ( X -S ⋆ ; a ) enjoys the following parameter-dependent upper and lower bounds:

Thus, an upper and lower bound to J ( X -S ⋆ ; a, π ) can be given by

In order to effectively tackle these bounds, we then study the properties of r J ( · ; r, π ) next.

Proposition E.5. Define

The function r J ( X -S ⋆ ; r, π ) with r ∈ [0 , 1] defined in (E.31) satisfies the following properties:

r J ( X -S ⋆ ; r, π ) is convex in r .

Proof. See §E.4.1 for the proof.

Upper and Lower Bounding J ( X -S ⋆ ; a, π ) . Previously, we show via a reformulation of f 5 that it suffices to bound J ( X -S ⋆ ; a, π ) . In the sequel, we let

It can be noticed that

where the second inequality follows from noting that the χ 2 -divergence defined as D χ 2 ( µ ∥ ν ) = ∑ x ( µ ( x ) -ν ( x )) 2 /ν ( x ) , and D χ 2 ( µ ∥ ν ) /D χ 2 ( ν ∥ µ ) ≤ sup x µ ( x ) /ν ( x ) .

Apparently, r π ( X -S ⋆ ; a ) is a function of a and enjoys the following parameter-dependent upper and lower bounds:

If a is small, we see that both r + ( a ) and r -( a ) are close to 1 , and we directly have

This suggests an upper bound of J ( X -S ⋆ ; a, π ) as

where the second line follows from the Lipschitz continuity property, and the last line holds because the ratio D + ( X -S ⋆ , π ) /D χ 2 ( µ π ( · | X -S ⋆ ) ∥ µ π ( · )) is upper bounded by ρ , and further by γ -1 . A similar lower bound can be obtained by changing the sign of γ -2 · max X -S ⋆ ,π µ π ( X -S ⋆ ) · ( e a -1) . Hence, we h

On the other hand, when a becomes large, we have both r + ( a ) and r -( a ) close to 0 , and we have

In a similar fashion, we have the following upper bound:

We can similarly obtain a lower bound by changing the sign of the second term inside the bracket. Hence, we have

Divergence of a . Recall that we have shown the dynamics of a in (E.30), where ξ ( a ) is negligible when L goes to infinity. Thus, when L is sufficiently large, we see by the nonnegativity of f 5 that a ( t ) continues to increase as t increases until it reaches a point where f 5 no longer dominates the approximation error. To characterize the regime where f 5 ≥ ξ ( a ) , we first note that for a ≤ log L it holds by (E.29) that

where ≈ hides logarithmic factors. For f 5 , we recall from Proposition E.4 that

where for small a we have f 5 = Ω(1) and for large a we have f 5 = Ω( e -2 a ) . Thus, e -2 a ≥ L -1 / 2 gives the condition for f 5 to dominate the approximation error, which gives a = O (log L ) . In the sequel, we consider the dynamics for a ≤ (log L ) / 8 and give a more rigorous analysis.

We use the notation x = o (1) to denote that a term is much smaller than 1 , for example, (log log L ) -1 = o (1) . For any x 0 and δ , we write x = x 0 ± δ to indicate that x is bounded within [ x 0 -δ, x 0 + δ ] . In the following, we assume there exists δ satisfying δ ≤ γ 2 / 4 ∧ 1 / 8 and

Note that

By additionally noting that µ π ( X -S ⋆ ) ≥ γ |S ⋆ | thanks to the lower bound of the transition probability, we are able to find such a δ if we have

where κ is defined as

and Ω( · ) only hides universal constants. Note that this is already guaranteed by the condition on L in (E.2). In particular, we can just take δ = γ 2 / 4 ∧ 1 / 8 in the following analysis.

Small a . Consider the case where a is small in the sense that µ π ( X -S ⋆ ) · ( e a -1) ≤ δ for any X -S ⋆ and π ∈ supp( P ) . In fact, one can directly deduce from our previous results that 1 -δ ≤ r -( a ) ≤ r π ( X -S ⋆ ; a ) < 1 and

For J ( X -S ⋆ ; a, π ) , we combine the condition that µ π ( X -S ⋆ ) · ( e a -1) ≤ δ with (E.32) to obtain that

Combining the above two results with Proposition E.4, we have

Also, the noise term ξ + ψ ( a ) is upper bounded by

by the construction of δ . Combining all the above results, we have the dynamics of a as

A simple reformulation gives

which implies that for small a , the growth follows

Therefore, in the beginning, a ( t ) grows super exponentially fast.

Large a . As a grows large such that µ π ( X -S ⋆ )( e a -1) ≥ δ -1 for all X -S ⋆ and π ∈ supp( P ) , we conclude that 0 < r π ( X -S ⋆ ; a ) ≤ r + ( a ) ≤ δ and

which imples that

For J ( X -S ⋆ ; a, π ) , we combine the condition that µ π ( X -S ⋆ ) · ( e a -1) ≥ δ -1 with (E.33) to obtain that

Combining the above two results with Proposition E.4, we have

For the noise term ξ + ψ ( a ) , we have

which can be verified by the condition on δ as well as the fact that we are only considering a ≤ (log L ) / 8 . We thus have for the gradient that

By rearranging the terms, we further have

Suppose this large a regime starts at t 0 with value a ( t 0 ) . Thus, for large a , the growth rate is characterized by

which is logarithmically fast. This step ends until a 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 f 5 . By the definition of f 5 in (E.28), we can rewrite f 5 as

where in the last step, we use the simple fact

In the definition of f 5 , the key quantity we aim to understand is the reweighted distribution r µ π ( z, Z | X -S ⋆ ) . For the readers' convenience, we copy the definition of the reweighted distribution here:

A key observation is that the reweighting only depends on the value of Z -S ⋆ . Let s S ⋆ = [ M ] \S ⋆ and denote by Z -s S ⋆ = ( z -h ) h ∈ s S ⋆ . Following the above observation, we can additionally condition on Z -S ⋆ and conclude that

as when fixing Z -S ⋆ , the exponential reweighting terms cancel out in the numerator and denominator in the definition (E.34). Using the above identity, we are able to expand r µ π ( z | X -S ⋆ ) as

where the second equality follows from the fact that the reweighing term in r µ π lifts the likelihood of Z -S ⋆ = X -S ⋆ by a factor of e a relative to the base distribution µ π ( Z -S ⋆ ) , and the denominator is just the normalization constant. In the sequel, we let r π ( X -S ⋆ ; a ) = (1 + µ π ( X -S ⋆ ) · ( e a -1)) -1 be the inverse of the normalization constant. We then have

On the other hand, by definition of r µ π in (E.34), we directly have

Combining both (E.35) and (E.36) we have for f 5 that

Here, we note that J ( · ; a, π ) is a function depending on both a and π , and can be expanded as

J ( X -S ⋆ ; a, π ) = ∑ k ∈ [ d ] ( µ π ( x = e k | X -S ⋆ ) -µ π ( x = e k )) 2 (1 -r π ( X -S ⋆ ; a )) µ π ( x = e k | X -S ⋆ ) + r π ( X -S ⋆ ; a ) µ π ( x = e k ) . Hence, we complete the proof of Proposition E.4.

Proof of Proposition E.5. Also, note that r J ( X -S ⋆ ; r, π ) is convex in r , as by taking the derivative of r J ( X -S ⋆ ; r, π ) with respect to r , we have

Hence, a naive upper bound for r J ( X -S ⋆ ; r, π ) is

where we remind the readers that D χ 2 ( µ ∥ ν ) = ∑ x ( µ ( x ) -ν ( x )) 2 /ν ( x ) . Next, we show that r J ( X -S ⋆ ; r, π ) is Lipschitz continuous in r :

where we use both the upper bound for r J ( X -S ⋆ ; r, π ) and the lower bound for the transition kernel that both µ π ( · | X -S ⋆ ) and µ π ( · ) are bounded between γ and 1 .

E.5 Lemma on GIH Approximation Error

Now given the convergence result for the training dynamics, the natural question to ask is how well the learned model implements the GIH mechanism. In the following part of this section, we state the lemma on the approximation error and also present a 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

where t 1 and t 2 are the ending time for the first two stages of the training, respectively. Suppose the error ∆ 1 , ∆ 2 = O ( L -1 ) after the first two stages' training, and a = Θ(log L ) after the last stage's training. Let y be the output of the model in (2.5) after the training and y ⋆ be the output of the GIH mechanism GIH ( x 1: L ; M,D ) defined in Definition 3.2. Then for any π ∈ supp( P ) and with high probability 1 -O ( L -1 ) , it holds that

Proof of Lemma E.6. Let s ⋆ l = ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) and s l = ⟨ u L +1 , u l ⟩ . Invoking Lemma F.1, the model misspecification error is bounded by

We note that the second layer's attention weight a can be as large as (log L ) / 8 . We are comparing the output of the model with the GIH mechanism GIH ( x 1: L ; M,D ) . Let N = ∑ l>M ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) . The output of this GIH mechanism is given by

We define

with σ ⋆ = ( σ ⋆ l ) l>M . Since ∥ x l ∥ 1 = 1 , the ℓ -1 norm of the difference between y ⋆ and the model's actual output is given by

Let us define the set Γ = { L ≥ l > M : ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) = 1 } and s Γ = { L ≥ l > M : ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) = 0 } . Using (E.37), for l ∈ Γ , we have 1 ≥ s l ≥ s ⋆ l -∆ = 1 -∆ and for l ∈ s Γ , we have 0 ≤ s l ≤ s ⋆ l +∆=∆ . Consider the normalization factor in the softmax function.

By the split of the set Γ and s Γ and noting that | Γ | = N , the normalization factor is lower and upper bounded by

Let us consider the event N ≥ 1 in the following. We then have for l ∈ Γ that

Note that a ∆ = o (1) due to the assumption that ∆ = O ( L -1 ) and a = o ( L ) . The right hand side is upper bounded by O ( a ∆ /N ) + O ( L exp( -a ) /N 2 ) . For l ∈ s Γ , we have

In summary,

The above inequality holds whenever N ≥ 1 . Now we aim to upper bound the probability that N = 0 . Note that N = ∑ L l = M +1 1 ( X l -S ⋆ = X L +1 -S ⋆ ) . We consider the following second moment:

where the first inequality holds by noting that D χ 2 ( µ ∥ ν ) = ∑ x ( µ ( x ) -ν ( x )) 2 /ν ( x ) and the last inequality holds by Lemma F.18. Therefore, by the Chebyshev's inequality, we have

We can take t = min E ∈X |S ⋆ | µ π ( E ) / 2 and by also taking a union bound over E ∈ X |S ⋆ | (which gives a d |S ⋆ | factor), we conclude that with high probability r O (1 -L -1 ) it holds that N ≥ tL = L · min E ∈X |S ⋆ | µ π ( E ) / 2 . Thus, it follows from (E.38) that with high probability

Hence, we complete the proof of Lemma E.6.

F Auxiliary Lemmas and Their Proofs

In this appendix, we present the auxiliary lemmas used to derive the approximation of the gradient flow dynamics in the proof of Theorem 3.6, which is presented in the previous appendix. The proofs of these lemmas are presented right below their statements.

F.1 Useful Inequalities

The following lemma provides a 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 u L +1 be the output feature after the FFN & Normalization layer. Then, the model misspecification error defined as

is bounded by ∆ 1 +∆ 2 , where ∆ 1 and ∆ 2 are the errors after the training of the first and second stages, respectively, and are defined respectively as

Proof of Lemma F.1. By definition of the output feature u l after the FFN & Normalization layer:

As each v ( h ) l is a convex combination of X M ( l ) where M ( l ) = { l -M,...,l -1 } , ∥ v ( h ) l ∥ 2 ≤ 1 . Thus,

where ∆ 1 is the error after the training of the first stage. Since v ( h ) l = ∑ j ∈ M σ ( h ) -j x l -j , we have

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 { ( i h , j h ) } h ∈S ⋆ arbitrarily, with each i h , j h ∈ [ H ] . Then we evaluate the product ∏ h ∈S ⋆ σ ( h ) -i h σ ( h ) -j h 1 ( x l -i h = x L +1 -j h ) given these indices. Then we sum over all possible values that { ( i h , j h ) } h ∈S ⋆ can take.

The above equation implies that

̸

̸

where the last inequality follows from the fact that

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 { ( i h , j h ) } h ∈S ∗ can take, except for the only case where ( i h , j h ) = ( h, h ) for all h ∈ [ H ] .

In summary, by triangle inequality, we have shown that

The proof is completed.

Next, in Lemma F.2, we establish a uniform bound for the quantity involved in the gradient.

Lemma F.2. Let y ( k ) = ∑ L l = M +1 σ l 1 ( x l = e k ) for each k ∈ [ d ] where ∑ L l = M +1 σ l = 1 and σ l ≥ 0 for all l ∈ [ L ] . Let ε and C be two positive numbers. For any C -bounded function f : X L +1 → [ -C, C ] , we have

Proof of Lemma F.2. By the triangular inequality, we have

where in the equality, we use the definition y ( k ) = ∑ L l = M +1 σ l 1 ( x l = e k ) and ∑ L l = M +1 σ l = 1 . 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 σ l ≈ 1 /L in the transformer model. The approximation error will be for g 0 , S to g 1 , S for Stage I and g h, 1 to g h, 2 for Stage II.

Lemma F.3. For the transformer model defined in (2.5) and any bounded function f : X L +1 → R such that sup x ∈X L | f ( x ) | ≤ C for a constant C > 0 , define two quantities A and B as

where s = u ⊤ L +1 U ⊤ 1: L and s y = ( L -M ) -1 ∑ L l = M +1 x l . Then, for all a ∈ [0 , 1] and ε ∈ (0 , 1] , it holds that

Proof of Lemma F.3. By triangular inequality, we have

Note that 0 ≤ s l ≤ 1 for all l = M + 1 , . . . , L thanks to the layer normalization. Then, for the softmax operation, we have

which implies that

Since indicator functions are bounded above by 1 , we have

For the second term, we have

where the last inequality follows from (F.2). Similarly, the following bound can be derived:

Combining (F.2), (F.3), (F.4) and (F.5), it holds that

where the last inequality follows from exp( x ) -1 ≤ 2 x for 0 ≤ x ≤ 1 . This concludes the proof of the lemma.

Lemma F.4 provides the approximation error introduced by µ π ( e k ) ≈ s y ( k ) in the transformer model. Lemma F.4. For the transformer model defined in (2.5) and any bounded function f : X L → R such that sup x ∈X L | f ( x ) | ≤ C for a constant C > 0 , define two quantities A and B as

where s y = ( L -M ) -1 ∑ L l = M +1 x l . Under Assumption 3.5, it holds that

where µ 0 ( · ) is the initial distribution over the first r n tokens X 1: r n . Here we let D χ 2 ( µ 0 ∥ µ π ) to denote D χ 2 ( µ 0 ( X 1: r n = · ) ∥ µ π ( X 1: r n = · )) , i.e., the χ 2 -divergence between µ n and the distribution over the first r n tokens under the stationary distribution µ π .

Proof of Lemma F.4. Let us use s y X ( · ) to remind the readers that s y ( · ) is also a function of X . We simplify the expectation E X | π by E in this proof. By rearranging the terms, we have

Here, we have three terms to control. For the first error term, we define

The first inequality above holds by noting that sup X | f ( X ) | ≤ C and the last inequality holds by noting that s y X ( e k ) = ( L -M ) -1 ∑ L l = M +1 1 ( x l = e k ) . Using Cauchy-Schwarz inequality, we arrive at

For the second term, we similarly have

Lastly, we have the error term

Here, the first term is upper bounded by Cγ -1 ε , and for the second term we have by Cauchy-Schwartz that

which shares a similar upper bound as err 1 . Now we invoke Lemma F.18 to conclude that

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 ∈ [ H ] ≤ D be a fixed set. For any h ∈ S , let r σ ( h ) and σ ( h ) be two fixed probability distributions over [ M ] . That is, for any i, j ∈ [ M ] , we have r σ ( h ) -i , σ ( h ) -j ∈ [0 , 1] , and ∑ M i =1 r σ ( h ) -i = ∑ M j =1 σ ( h ) -j = 1 . Given these distributions over [ M ] , we define

where we let x l ∈ X denote the l -th token in the Markov chain for all l ∈ [ L +1] . Moreover, with slight abuse of notation, we let ( z, Z ) = ( z, z -1 , . . . , z -M ) ∈ X M +1 and ( x, X ) = ( x, x -1 , . . . , x -M ) ∈ X M +1 be two independent random variables sampled from the stationary distribution µ π . We define random variables r v ( h ) ( Z ) and v ( h ) ( X ) as

Using r v ( h ) L +1 , v ( h ) l , r v ( h ) ( Z ) , and v ( h ) ( X ) , we define two quantities A and B as

where E X | π means that the expectation is taken with respect to the randomness of the Markov chain with transition π . Then, under Assumption 3.5, we have

where µ 0 ( · ) is the initial distribution over the first r n tokens X 1: r n and D χ 2 ( µ 0 ∥ µ π ) is a short-hand notation of D χ 2 ( µ 0 ( X 1: r n = · ) ∥ µ π ( X 1: r n = · )) .

Proof of Lemma F.5. By triangular inequality, we have

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 p π ( X ) denote the joint distribution of the whole sequence X under kernel π . By the definitions of r v ( h ) L +1 and v ( h ) l , we have

where we use ( i h , j h ) as the indices to highlight that they are associated with head h . And we use k h ∈ [ d ] to index all the possible common values for x l -i h and x L +1 -j h . Then plugging this equality into ∏ h ∈S ⟨ v ( h ) l , v ( h ) L +1 ⟩ and exchanging the order of product and summation, we have

where the summation means that we sum over all possible values that { i h , j h , k h } h ∈S and k can take. Specifically, each i h and j h take values in [ M ] , and each k h and k takes values in [ d ] . Moreover, using the property of indicator functions, we can further simplify (F.6) by gathering all indicators:

Now we take expectations with respect to the randomness of X on both ends of (F.7) and get

To further simplify the above equality, we define a new probability distribution over X L +1 -M : L +1 and another subsequence of length M +1 . Note that X L +1 -M : L +1 contains is a subsequence with M + 1 tokens. We let ( z, Z ) = ( z, z -1 , . . . , z -1 , z -M ) denote a random token sequence of size M +1 in reverse order. We define a joint distribution p p π over X L +1 -M : L +1 and ( z, Z ) as follows. Let E = ( E 0 , E -1 , . . . , E -M ) and E ′ = ( E ′ 0 , E ′ -1 , . . . , E ′ -M ) be two elements in X M +1 . That is,

each component of E and E ′ are in X . The probability mass function of p p π is defined as

That is, p p π can be viewed as the joint distribution of X L +1 -M : L +1 with an averaged distribution of the history. When L is sufficiently large, by the mixing property of the Markov chain, we expect that, under p p π , ( z, Z ) is approximately independent of X L +1 -M : L +1 , and the marginal distributions of ( z, Z ) and X L +1 -M : L +1 are both close to the stationary distribution µ π . We will translate this intuition into a rigorous argument in Lemma F.17, which bounds the total-variation distance between p p π and the product distribution µ π × µ π .

With p p π defined in (F.8), we can rewrite the expectation above as

Similarly, by the definitions of r v ( h ) ( Z ) and v ( h ) ( Z ) , we can write ⟨ r v ( h ) ( Z ) , v ( h ) ( X ) ⟩ as

Then, multiplying these terms with h ∈ S , we can write

Recall that here ( z, Z ) = ( z, z -1 , . . . , z -M ) and ( x, X ) = ( x, x -1 , . . . , x -M ) are independently sampled from the stationary distribution µ π . Taking the expectation under µ π , we have

To bound the first absolute value term in the upper bound on | A -B | , we aim to compare (F.9) and (F.11). To this end, let us fix collections of index pairs ( i h , j h ) h ∈S . Let S 1 = { i h : h ∈ S} and S 2 = { j h : h ∈ S} be the unique values in ( i h ) h ∈S and ( j h ) h ∈S . Since there might exists two elements h and h ′ in S such that i h = i h ′ or j h = j h ′ , |S 1 | and |S 2 | might be strictly less than |S| . As a result, p p π ( x L +1 = z = e k , x L +1 -i h = z -j h = e k h , ∀ h ∈ S ) only involves random variables x L +1 , X L +1 -S 1 = { x L +1 -i } i ∈S 1 , z , Z -S 2 = { z -j } j ∈S 2 , which are a subset of the random variables defined in (F.8). Similarly,

only involves a subset of random variables x , X -S 1 = { x -i } i ∈S 1 , z , and Z -S 2 . Let us define s E = ( E 0 , ( E -i ) i ∈S 1 ) ∈ X |S 1 | +1 and s E ′ = ( E ′ 0 , ( E ′ -j ) j ∈S 2 ) ∈ X |S 2 | +1 . By enumerating s E

in X |S 1 | +1 and s E ′ in X |S 2 | +1 , we equivalently enumerate all possible values the above random variables can take. Therefore, by comparing (F.9) with (F.11), we have

where in the last line, we use Y and Y ′ as placeholders for the random variables ( x L +1 , X L +1 -S 1 ) and ( z, Z -S 2 ) respectively. In the first equality, we sum over s E ∈ X |S 1 | +1 and s E ′ ∈ X |S 2 | +1 , 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

where the last inequality holds by Corollary F.15 and the fact that |S 1 | ≤ |S| . Specifically, Corollary F.15 implies that the density function of the joint distribution of x L +1 and X L +1 -S 1 is lower bounded by γ |S 1 | +1 ≥ γ |S| +1 . Thus, combining (F.12) and (F.13), we have

Therefore, to bound the first absolute value term, we combine (F.9), (F.11), and (F.14) and use triangle inequality to get

Furthermore, recall that r σ ( h ) and σ ( h ) are probability distributions over [ M ] for all h ∈ S . By going over all possible values that { ( i h , j h ) } h ∈S can take, we have

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).

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 ⟨ v ( h ) l , v ( h ) L +1 and ⟨ r v ( h ) ( Z ) , v ( h ) ( Z ) ⟩ and do not have indicators 1 ( x L +1 = x l ) and 1 ( x = z ) .

Similar to the derivation in (F.9) and (F.11),

Similar to (F.12), for any fixed collection of index pairs ( i h , j h ) h ∈S , we let S 1 = { i h : h ∈ S} and S 2 = { j h : h ∈ S} denote the unique values in ( i h ) h ∈S and ( j h ) h ∈S . By Lemma F.17, we have

Here we use r Y and r Y ′ as placeholders for random variables X L +1 -S 1 and Z -S 2 . We note that Lemma F.17 can be applied to any subsets of X L +1 -M : L +1 and ( z, Z ) . Therefore, combining (F.16), (F.17), and (F.18), we conclude that

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 χ 2 -mutual information.

Lemma F.6. Consider a fixed set S ∈ [ H ] ≤ D . For any h ∈ S , let r σ ( h ) and σ ( h ) be two probability distributions over [ M ] . That is, for any i, j ∈ [ M ] , we have r σ ( h ) -i , σ ( h ) -j ∈ [0 , 1] , and ∑ M i =1 r σ ( h ) -i = ∑ M j =1 σ ( h ) -j = 1 . Moreover, we let ( z, Z ) = ( z, z -M , . . . , z -1 ) ∈ X M +1 and ( x, X ) = ( x, x -M , . . . , x -1 ) ∈ X M +1 be two independent random variables sampled from the stationary distribution µ π . We define random variables r v ( h ) ( Z ) and v ( h ) ( X ) as

Let ( i ⋆ h , j ⋆ h ) h ∈S be any fixed collection of index pairs, where i ⋆ h ∈ [ M ] and j ⋆ h ∈ [ M ] for all h ∈ S . We define quantities A and B as

Under Assumption 3.5, it holds that

Proof of Lemma F.6. To simplify the notation, we define a signal set Γ( S ) and an error set s Γ( S ) as

Similar to (F.10), we can write E π [ A ] as

where we exchange the order of product and summation. Using the notation s Γ( S ) , we can split the summation into two parts:

Here last equality holds by the definition of the B and the fact that r σ ( h ) and σ ( h ) are fixed vectors.

Therefore, to prove this lemma, it suffices to upper bound the second term above. To this end, we apply Lemma F.7 stated below for any fixed set of indices ( i h , j h ) h ∈S ∈ s Γ( S ) . Specifically, let S 1 = { i h : h ∈ S} and S 2 = { j h : h ∈ S} denote the unique values of ( i h ) h ∈S and ( j h ) h ∈S . Lemma F.7 implies that

Combining (F.19) with the fact that

the desired term is bounded above by (1 -∏ h ∈S σ ( h ) -i ⋆ h r σ ( h ) -j ⋆ h ) · I χ 2 ( S ⋆ ) , which concludes the proof.

Lemma F.7. Let S ∈ [ H ] ≤ D be a fixed subset and let { ( i h , j h ) } h ∈S be a fixed collection of index pairs, where i h , j h ∈ [ M ] for all h ∈ S . Let S 1 = { i h : h ∈ S} and S 2 = { j h : h ∈ S} denote the unique values of ( i h ) h ∈S and ( j h ) h ∈S . We let ( z, Z ) = ( z, z -M , . . . , z -1 ) ∈ X M +1 and ( x, X ) = ( x, x -M , . . . , x -1 ) ∈ X M +1 be two independent random variables sampled from the stationary distribution µ π , where π is the transition kernel of the Markov chain and is sampled from prior P . If Assumption 3.5 holds, it follows that

where r I χ 2 ( S ) is the modified χ 2 -mutual information defined in Definition 3.1 and S ⋆ = argmax S∈ [ H ] ≤ D r I χ 2 ( S ) .

̸

Proof of Lemma F.7. We first note that it is allowed |S 1 | = |S 2 | as there could be duplicate values in both ( i h ) h ∈S and ( j h ) h ∈S , while S 1 and S 2 are the unique values. In the sequel, we let X -S 1 denote { x -i h } h ∈S and let Z -S 2 denote { z -j h } h ∈S , where repeated elements are removed. Moreover, we let

{ X -S 1 = Z -S 2 } be the event that x -i h = z -j h for all h ∈ S . Notice that ∏ h ∈S 1 ( x -i h = z -j h ) = 1 ( X -S 1 = Z -S 2 ) . Then, we have

Here in the second equality, we take a conditional expectation given X -S 1 and Z -S 2 . The last equality can be verified by direct computation. To simplify the expectation above, we aim to transform the indicator of 1 ( X -S 1 = Z -S 2 ) into probabilities involving X -S 1 and Z -S 2 . To this end, we need to explicitly enumerate all possible values that X -S 1 and Z -S 2 can take. This is challenging, as there may be duplicated values in both i h and j h , and thus X -S 1 and Z -S 2 can have different sizes. However, since S 1 is a 'reduction' of { i h } h ∈S , we can revert to the original space and consider E = ( E h ) h ∈S ∈ X |S| that respects the reduction from { i h } h ∈S to S 1 . Here each E h is the value x i h takes. In other words, with ( i h ) h ∈S that might have duplicated values, we consider the values taken by ( x i h ) h ∈S , with duplicates allowed. And E has the same duplication structure as ( x i h ) h ∈S . In the following, we describe these values by introducing the notion of compatibility.

̸

Definition F.8 (Compatible Value Set) . We say that E ∈ X |S| is compatible with ( i h ) h ∈S if, for any h = h ′ such that i h = i h ′ , we have E h = E h ′ . In other words, the unique values in E can be indexed by { i h } h ∈S = S 1 if E is compatible with ( i h ) h ∈S .

By this definition, E is compatible with ( i h ) h ∈S if it respect duplication pattern of ( x i h ) h ∈S . If i h = i h ′ , then we know that x i h and x i h ′ is the same token. Since x i h and x i h ′ take values E h and E h ′ , we must have E h = E h ′ . As a concrete example, suppose S = { 1 , 2 , 3 } , and the values of ( i h ) h ∈S are given by ( i 1 , i 2 , i 3 ) = (1 , 2 , 1) . Therefore, we have S 1 = { 1 , 2 } , which contains the unique values of ( i 1 , i 2 , i 3 ) . Now, let E = ( E 1 , E 2 , E 3 ) . For E to be compatible with ( i h ) h ∈S , we must have E 1 = E 3 since i 1 = i 3 . There is no restriction on E 2 . So, a compatible value set for this example could be E = ( a, b, a ) , where a and b are elements of X .

In the sequel, we define E as the set of vectors in X |S| that are compatible with both { i h } h ∈S and { j h } h ∈S , i.e.,

The compatibility condition allows us to assign x -i h , z -j h the value E h for all h ∈ S when E ∈ E . Under this assignment, the constraint X -S 1 = Z -S 2 is automatically satisfied. We use the notation { X -S 1 = E } to denote the event that x -i h = E h for all h ∈ S , and similarly for Z -S 2 = E . In particular, we are able to rewrite the indicator 1 ( X -S 1 = Z -S 2 ) as ∑ E ∈E 1 ( X -S 1 = E,Z -S 2 =

E ) . Then we can rewrite (F.20) by separating X -S 1 and Z -S 2 as

where in the last inequality, we apply ab ≤ a 2 + b 2 / 2 .

Next, for each E ∈ E , consider E ′ = ( E ′ i ) i ∈S 1 such that

Note that for each 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 1 . The uniqueness is due to the fact that (F.22) completely determines all the values in E ′ because enumerating over i h for h ∈ S is just the same as enumerating over i for i ∈ S 1 . In fact, E ′ contains all the unique values of E . In the above example, we have S 1 = { 1 , 2 } and thus E ′ = ( a, b ) when E = ( a, b, a ) .

Since E ′ is uniquely defined based on E , we are able to define an operator J 1 that maps E ∈ E to E ′ ∈ X |S 1 | according to the mapping given in (F.22). Let J 1 ( E ) be the image of E under J 1 . It is important to note that for each E ′ ∈ J 1 ( E ) , there is also a unique pre-image E ∈ E such that J 1 ( E ) = E ′ according to the rule (F.22). Therefore, J 1 is an one-to-one mapping from E to J 1 ( E ) . In the following, for any E ′ ∈ J 1 ( E ) , we denote by { X -S 1 = E ′ } the event where x -i = E ′ i for all i ∈ S 1 . Equivalently, we have x -i h = E ′ i h = E h for all h ∈ S . Thus, the event { X -S 1 = E ′ } is exactly the same as { X -S 1 = E } introduced above. Therefore, the first term on the right hand side of (F.21) can be reformulated as

where the equality follows from the bijection between E and J 1 ( E ) , and the last inequality holds by noting that J 1 ( E ) ⊆ X |S 1 | . 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

Lastly, note that r I χ 2 ( S ) ≤ r I χ 2 ( S ⋆ ) for any S ∈ [ H ] ≤ D by the optimality of S ⋆ . Hence, we complete the proof of Lemma F.7.

Lemma F.9 quantifies the approximation error from σ l ≈ σ ⋆ l and y ( k ) ≈ y ⋆ ( k ) for Stage III.

Lemma F.9. For the transformer model defined in (2.5) , define two quantities f 1 and f 2 as

where the expectation is taken over all the randomness in the data, and

with S ⋆ is the optimal information set. Under Assumption 3.5, it holds that

where ∆ 1 :=1 -p S ⋆ and ∆ 2 :=1 -∏ h ∈S ⋆ ( σ ( h ) -h ) 2 .

Proof. We separate the approximation error into three parts | f 1 -f 2 | ≤ err 1 +err 2 +err 3 , which are explained in detail as follows.

The First Error Term. Here, the first error err 1 captures the error of replacing ∏ h ∈S ⋆ 1 ( x l -h = x L +1 -h ) with ∏ h ∈S ⋆ ⟨ v ( h ) l , v ( h ) L +1 ⟩ in f 2 :

Then using Lemma F.1, we conclude that

The Second Error Term. The second error term characterizes the difference in σ l and σ ⋆ l :

To characterize such an error, we invoke equation (53) in Lemma 5.1 of Chen et al. (2022). This lemma states that for σ and σ ⋆ being the output of the softmax function with scaling parameters a for s and s ⋆ respectively, i.e.,

it holds that ∥ σ -σ ⋆ ∥ 1 ≤ 4 a · ∥ s -s ⋆ ∥ ∞ . Consequently, we have ∥ σ -σ ⋆ ∥ 1 ≤ 4 a · (∆ 1 +∆ 2 ) by Lemma F.1. We notice that

Thus, err 2 ≤ 4 aε -1 · (∆ 1 +∆ 2 ) .

The Third Error Term. The last error term characterizes the difference between y ⋆ and y :

By noting that y ⋆ = ∑ L l = M +1 σ ⋆ l x l and y = ∑ L l = M +1 σ l x l , we have

The first term of err 3 can be bounded by

Moreover, for the second term of err 3 , we have

It then holds that

Therefore, we complete the proof of Lemma F.9.

Lemma F.10. Let us define for brevity,

where Z = ( z -M , . . . , z -1 ) and µ π is the stationary distribution of the Markov chain over a window of size M +1 . We denote by r µ π X ( e k ) = r µ π X ( z = e k ) where r µ π X ( z ) is the marginal distribution for z and serves as the population counterpart for y ⋆ = ∑ L l = M +1 σ ⋆ l x l . We define quantity A and B as

Under Assumption 3.5, we have

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

To handle this error, we define three error terms as

For the first error term, we have that

where we recall that by assumption, γ provides a lower bound for π ( · | X pa ) , hence also a lower bound for r µ π X ( e k ) . Next, we invoke Proposition F.19 which provides an upper bound for the difference between the empirical and population distributions in terms of the ℓ 1 -norm:

Hence, we control the first error term.

For the second error term, we follow the same procedure and obtain an upper bound as

For the last error term, it holds that

where the last inequality follows directly from (F.23).

In summary, the difference between f 2 and f 3 is bounded by

| f 2 -f 3 | ≤ err 1 +err 2 +err 3 ≤ 8(1 -λ ) -1 / 2 ( D χ 2 ( µ 0 ∥ µ π ) + 1) 1 / 4 +8 √ M L 1 / 2 · γ |S ⋆ | +1 + 2 dε γ , which completes our proof of Lemma F.10.

The following lemmas are for analyzing the error | f 3 -f 4 | for Stage III. Lemma F.11. We define

where

Under Assumption 3.5, we have

Proof of Lemma F.11. For Z = ( z -M , . . . , z -1 ) and Z ′ = ( z ′ -M , . . . , z ′ -1 ) , we let Z -S ⋆ = ( z -h ) h ∈S ⋆ , we define

Using these notations, we can rewrite the normalizing factor in r µ π X and σ ⋆ l respectively as

We also define

If we further define p ν π X ( z, Z -S ⋆ ) = ∑ L l = M +1 1 ( x l = z, X l -S ⋆ = Z -S ⋆ ) , then we have

Using the above definitions and relationship, A and B can be rewritten as

Therefore, the difference between A and B is given by

where the last inequality follows from the result in Proposition F.19. Invoking the lower bound µ π ( x L +1 , X L +1 -S ⋆ ) ≥ γ |S ⋆ | +1 , we complete the proof of Lemma F.11.

F.3 Lemmas on Concentration of Markov Chain

Recall that we previously define X = ( x 1 , . . . , x L ) as the observed sequence and x L +1 as the value at time L +1 to be predicted. For generality, we will use X = ( x 1 , . . . , x L +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 1 , . . . , -r n } , and as the start of a chain, we sample the first r n tokens by ( x 1 , . . . , x r n ) ∼ µ 0 .

In the sequel, we will study concentration properties of the Markov chain X for a window of tokens with window size at most M , where M > r n . To proceed, let us consider a fixed set S ⊆ [ M ] . For any l ∈ [ M + 1 , L + 1] , we define Y l = ( x l , X l -S ) as a new vector containing the token at position l and also the tokens in the past S positions prior to x l . Here, we follow the convention that X l -S = ( X l -i ) i ∈S . We also consider another fixed subset S ′ ⊆ [ M ] and similarly define Y ′ l = ( x l , X l -S ′ ) .

The concentration properties of the Markov chain are rooted in 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 = {-r 1 , . . . , -r n } , we aim to make Y L +1 approximately independent of Y l by conditioning on some intermediate parent sets. To this end, we define A = ( x L +1 -M , . . . , x L -M + r n ) ∈ X r n and B l = ( x l -r n +1 , . . . , x l ) ∈ X r n as these intermediate parent sets. By the Markov property and the parent set structure, we have the following conditional independence relations:

To illustrate, let us consider the first condition Y L +1 ⊥ ⊥ ( B l , Y l ) | A . When l ≤ L -M + r n , the B l and Y l are both contained in the history { x k : k ≤ L -M + r n } = A ∪ { x k : k ≤ L -M } . When conditioning on A , the randomness of ( B l , Y l ) is measurable by the σ -algebra generated by the 'past' { x k : k ≤ L -M } . Moreover, the randomness of Y L +1 is measurable by the σ -algebra generated by the 'future' { x k : k ∈ [ L +1 -M + r n , L +1] } when conditioning on A . Notice that the parent to the any element in the future { x k : k ∈ [ L +1 -M + r n , L +1] } is either contained in A , or can be generated conditioned on A without touching further history { x k : k ≤ L -M } . Thus, by the Markov property, conditioning on A , Y L +1 is independent of the past { x k : k ≤ L -M } , and in particular, ( B l , Y l ) . Similarly, since B contains the parent of x l +1 , conditioning on B , Y l is independent of x l +1 and later tokens. Moreover, given B , the randomness of Y l comes from the randomness of x l -M , . . . , x l -r n . Since l ≤ L -M + r n , we have L +1 -M ≥ l +1 -r n . As a result, conditioning on B , the randomness of ( Y L +1 , A ) comes from tokens generated no earlier than x l +1 . Therefore, ( Y L +1 , A ) and Y l are conditionally independent given B l . We visualize the definition of Y L +1 , A , B l , and Y l in Figure 10

2 YL+1

2 Y

XL+1

XL+1-M+rn

XL-M+rn

XL+1-M

x1

X1-rn+1

x1-rn

X1-M

A

B

Similarly, for Y ′ l = ( x l , X l -S ′ ) defined using the subset S ′ , we also parallel conditional independence relations:

In particular, we also have

Using { Y l , Y ′ l } , we define a joint distribution p p π over 2 + |S| + |S ′ | tokens as follows. For any E ∈ X |S| +1 and E ′ ∈ X |S ′ | +1 , the probability mass function of p p π is defined as

Here, Y ′ is just a placeholder for Y ′ l as p p takes an average over l and does not depend on any specific position index. The summation ∑ A,B l means we sum over all possible values that A and B l can take. In the last line of (F.25), we decompose the joint distribution p π ( Y L +1 = E,Y ′ l = E ′ ) into the product of the conditional distributions by the Markov property in (F.24). That is,

Here the second equality follows from the fact that ( Y L +1 , A ) ⊥ ⊥ Y ′ l | B l and the last equality follows from the fact that Y L +1 ⊥ ⊥ ( B l , Y ′ l ) | A , which implies p π ( Y L +1 = E | A,B l ) = p π ( Y L +1 = E | A ) . Moreover, we denote by P i π 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

In the following, we always consider a fixed transition kernel π and omit the superscript/subscript π in the matrix notation. We denote the transition matrix by P π and the stationary distribution by µ π for a window of length r n . For the transition matrix, we index each row by the next r n -window Z ′ and each column by the current r n -window Z . Under this notation, since both A and B l have lengths r n , we have

Here P L -M + r n -l π ( A | B l ) corresponds to the ( A,B l ) -entry of the matrix ( P π ) L -M + r n -l . Combining (F.24) and (F.27), we obtain the last equality in (F.25).

In the sequel, to simplify the notation, we write P π and µ π as P and µ respectively. Let us consider the reweighted transition kernel

where √ µ is the element-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 √ µ , i.e. √ µ = K √ µ and √ µ ⊤ = √ µ ⊤ K . However, the projection in the leading eigenspace (or the Perron projection) is not of 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 i ≥ 0

Proof of Proposition F.12.

where the last equality holds by noting that K - √ µ √ µ ⊤ project √ µ to the zero vector, and for any v ⊥ √ µ , we have ( K - √ µ √ µ ⊤ ) v = Kv . Thus for any test vector x :

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 kernel π admits a lower bound γ > 0 , then we can guarantee that both p π and µ π admit a uniform lower bound.

Proposition F.13 (Uniform Lower Bound) . Suppose π ( · | X pa ) ≥ γ uniformly for some γ > 0 and pa = {-r 1 , . . . , -r n } . Suppose X 1: r n ∼ µ 0 ( · ) where µ 0 ∈ ∆( X r n ) . Then for any S tokens x l 1 , x l 2 , . . . , x l S such that l s ≥ r n for any s ∈ [ S ] , we have

Using Proposition F.13, we show that the transition matrix P π is primitive.

Corollary F.14 (Uniform Lower Bound Implies Primitive Transition) . Under the condition of Proposition F.13, with π ( · | X pa ) ≥ γ > 0 , the transition matrix defined in (F.26) is primitive.

Proof of Corollary F.14. If the initial distribution is set to be any one-hot vector in ∆( X r n ) , and taking x l 1 , . . . , x l S in Proposition F.13 to be x r n +1 , . . . , x 2 r n , we conclude that p π ( X r n +1:2 r n | X 1: r n ) > 0 holds for any X r n +1:2 r n , X 1: r n ∈ X r n . Recall from the definition that for a primitive matrix P , we can find some positive integer k such that P k has all positive entries. For our case, we can set k = r n and everything follows by noting that p π ( X r n +1:2 r n | X 1: r n ) = P r n π ( X r n +1:2 r n , X 1: r n ) .

Another corollary of Proposition F.13 is that, if we take µ 0 = µ π , which is the stationary distribution, we can replace p π in Proposition F.13 by µ π .

Corollary F.15. Suppose π ( · | X pa ) ≥ γ uniformly for some γ > 0 and pa = {-r 1 , . . . , -r n } . For the stationary distribution µ π and S tokens x l 1 , x l 2 , . . . , x l S such that l s ≥ r n for any s ∈ [ S ] , we have µ π ( x l 1 , . . . , x l S ) ≥ γ S .

We prove Proposition F.13 as follows.

Proof of Proposition F.13. Without loss of generality, suppose that M ≤ l 1 < l 2 < . . . < l S . We will prove the statement by induction on the number of tokens S . If S = 1 , we can rewrite

Now, suppose the statement holds for 1 , 2 , . . . , S -1 . Let Y = x l 1 , . . . , x l s -1 . Then, we have

where the last inequality holds by the induction condition. Hence, we finish the proof.

Before analyzing p p π , we first study a simpler convergence result: quantifying the closeness between ∑ L l = M +1 η L -l p π ( B l = b ) / ∑ L l = M +1 η L -l and µ π ( b ) for certain values of η ∈ (0 , 1] .

Lemma F.16. Following the notations introduced above, for the Markov chain with parent set pa = {-r 1 , . . . , -r n } , let D χ 2 ( µ 0 ∥ µ π ) be the χ 2 -divergence between the initial distribution µ 0 and the stationary distribution µ π over the first r n tokens. Take any S ⊆ [ M ] and let Y l = ( x l , X l -S ) for l = M +1 , . . . , L +1 . Suppose L/ 2 ≥ M ≥ r n . We have

Proof of Lemma F.16. Let c l = η L -l / ∑ L -M + r n l = r n η L -l , where η ∈ [0 , 1] is a constant to be determined. Denote by µ 0 , a vector of length |X| r n , the initial distribution of the chain. We begin by quantifying the total variation (TV) distance:

Let b ∈ X r n , representing the value for a lengthr n window. Using matrix notation, we have:

where 1 b is the indicator vector corresponding to b . The last equality follows from Proposition F.12. For any test vector u ∈ { 0 , 1 } |X| rn , using the variational representation of TV distance:

where the inequality holds by ∥ u ⊤ diag ( √ µ ) ∥ 2 ≤ ∥ √ µ ∥ 2 = 1 and K - √ µ √ µ ⊤ has leading eigenvalue with magnitude λ . The last identity follows directly from the definition of the χ 2 -divergence that D χ 2 ( µ 0 ∥ µ π ) + 1 = ∑ b µ 0 ( b ) 2 /µ π ( b ) .

Substituting the definition of c l , we have

We consider two special cases. In the first case, we set η = λ , which gives us

In the second case, we set η = 1 , which gives us

Note that the TV distance is an f -divergence. Thus, we can use the data processing inequality to obtain the desired result for Y l from the above inequality. To do so, note that ∑ L l = M +1 p π ( Y l = · ) / ( L -M ) and µ π ( Y L +1 = · ) can be transformed from ∑ L -M + r n -1 l = r n p π ( B l = · ) / ( L -M ) and µ π ( A = · ) by the same emission kernel

π π π π

Therefore, by the data processing inequality, it holds that

Similarly for p π ( Y L +1 = · ) and µ π ( · ) , we have

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 ∑ L l = M +1 p π ( Y l = · ) / ( L -M ) converges to µ π ( A = · ) in total variation distance. This represents a 'first-order' convergence since it involves the average of the marginal distribution of Y l . However, the quantity of interest in (F.25) is the average of the joint distribution of Y L +1 and Y l , 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 = {-r 1 , . . . , -r n } , let D χ 2 ( µ 0 ∥ µ π ) be the χ 2 -divergence between the initial distribution µ 0 and the stationary distribution µ π over the first r n tokens. Take any S , S ′ ⊆ [ M ] and let Y l = ( x l , X l -S ) and Y ′ l = ( x l , X l -S ′ ) for l = M +1 , . . . , L +1 . Suppose L/ 2 ≥ M ≥ r n . For p p π defined in (F.25) , we have

In particular, we have

Proof of Lemma F.17. Let us take µ π ( E ) · ( L -M ) -1 ∑ L l = M +1 p π ( Y l = E ′ ) as the intermediate distribution, and we have by (F.25) that

where we use the fact that ∑ A µ π ( Y L +1 = E | A ) µ π ( A ) = µ π ( Y = E ) for the first line. The second term on the right hand side can be easily controlled as we already have an L -1 factor. We let TV 0 be the total variation distance of the second term. It is easy to see that

where we remark that (II) is a function of both E and E ′ , and the total variation distance is just taking the sum of the absolute values of the differences. Here, we also use the fact that L ≥ 2 M . Using Proposition F.12, we can also rewrite the first term on the right hand side of (F.32) in the matrix form as

When considering the ℓ 1 -norm of the above term, we introduce a test matrix U of shape |X| | Y L +1 | × |X| | Y L +1 | with each element of U chosen from { 0 , 1 } . Let TV 1 be the total variation distance of the first term (I). Then, we have

To upper bound this quantity, we consider each row of U as U ( E, · ) = u ( · | E ) ⊤ . Note that u ( · | E ) is also a { 0 , 1 } -valued vector. By expanding the trace, we have

Note that the ℓ 2 -norm of the vector in the last line can be upper bounded by

where the first inequality holds by noting that p π ( Y ′ l = · | B l = · ) ⊤ · u ( · | E ) is a vector with element within [0 , 1] , and also invoking the operator norm of the matrix r K - √ µ √ µ ⊤ . The second identity follows from the definition of the χ 2 -divergence that D χ 2 ( p π ( B l = · ) ∥ µ π ( · )) + 1 = ∑ b p π ( B l = b ) 2 /µ π ( b ) . The last inequality is the data processing inequality as p π ( B l = · ) can be transformed from µ 0 ( B r n ) and µ π ( B l ) can be transformed from µ π ( B r n ) by the same emission kernel µ π ( B l = · | B r n = · ) . Consequently, we have for the TV distance that

where in the first equality, we use the variational form of the ℓ 2 -norm for vector µ π ( Y L +1 = E | A = · ) · diag( √ µ ) . In the second inequality, we apply (F.33) and use the Bayes rule. The last inequality follows from the Cauchy-Schwarz inequality. Here, the mutual information I χ 2 ( A ; Y L +1 ) + 1 can be upper bounded by

and the last term involving v ( A | E ) can be upper bounded by 1 thanks to the constraint on v ( · | E ) . In conclusion,

Lastly, let us relate the intermediate distribution to the final distribution µ π ( Y = · ) × µ π ( Y ′ = · ) , where we define the total variation distance TV 2 as

Invoking (F.28) of Lemma F.16, we have this quantity upper bounded by

Using the triangular inequality for the total variation distance, we have

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.

Lemma F.18. For the χ 2 -divergence between the empirical distribution ( L -M ) -1 ∑ L l = M +1 1 ( Y l = · ) and the stationary distribution µ π ( · ) , we have

where the expectation is with respect to X ∼ p .

Proof of Lemma F.18. By definition of the χ 2 -divergence, what we aim to bound is just

To study the above quantity, for l ≥ 2 M -r n +2 , we define

Following our convention, we let A l = X l -M : l -M + r n -1 and B l ′ = X l ′ -r n +1: l ′ be two lengthr n window and by the Markov property, we have

Let us fix an index l ≥ 2 M -r n +2 and take a summation over M +1 ≤ l ′ ≤ l -M + r n -1 . Expanding the joint distribution, we have

where the first identity follows from the fact that

and the second identity follows from Proposition F.12. We next invoke the Cauchy-Schwarz inequality for trace, i.e., Tr( W ⊤ V ) 2 ≤ Tr( W ⊤ W ) Tr( V ⊤ V ) , where we take

Note that

Following the same calculation, we have

Therefore,

We further have

where the last inequality holds by the data processing inequality. Similarly, we have

Therefore, we conclude that

and

where we double the value as l > l ′ only contributes to half of the terms in the double summation. Note that in the above summation for l > l ′ , we only include terms satisfying l -l ′ ≥ M -r n +1 and l -( M +1) ≥ M -r n +1 . For the remaining ( l, l ′ ) not included above, each term is bounded above by

and we have no more than 4 L ( M -r n +1) of these terms in total. As a result, we conclude with L/ 2 ≥ M ≥ r n that

Hence, we complete the proof of Lemma F.18.

Proposition F.19. Let us define

where Z -S ⋆ = ( z -h ) h ∈S ⋆ and µ π is the stationary distribution of the Markov chain over a window of size M +1 . We also treat r µ π X ( · ) as a length |X| vector where X is the state space of the Markov chain. Let p ν π X ( z, Z -S ⋆ ) = ∑ L l = M +1 σ ⋆ l 1 ( x l = z, X l -S ⋆ = Z -S ⋆ ) where

Then, we have

Proof of Proposition F.19. To unify the notations, we let Z = ( z -M , . . . , z -1 ) and define

Using these notations, we can define the normalizing factor in r µ π X and y ⋆ X respectively as

We also define

We can then rewrite the objective as

Furthermore, notice that

where we define Γ X = { Z -S ⋆ : Z -S ⋆ = X L +1 -S ⋆ } . Note that when Z -S ⋆ ∈ Γ X , we have R ( Z -S ⋆ , X L +1 -S ⋆ ) = e a and when Z -S ⋆ / ∈ Γ X , we have R ( Z -S ⋆ , X L +1 -S ⋆ ) = 1 . For the first term on the right-hand side of (F.34), we have by Cauchy-Schwarz that

where in the last inequality, we invoke Lemma F.18 where we take Y l = x l in the lemma. For the second term on the right hand of (F.34), we note that

where the last two inequalities follow from the Cauchy-Schwarz inequality. We have an upper bound for the second term on the right-hand side of (F.35) that

We can also apply Lemma F.18 to the first term with Y L +1 = ( x L +1 , X L +1 -S ⋆ ) and conclude that

In summary, we have

Note that the second term is dominated by the first term. Thus, we conclude the proof of Proposition F.19.

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 as long as it 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 no 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 in 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 on 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 experimental 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: [Yes]

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 paper reproducible is important, regardless of whether the code and data are provided or not.

If 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 a specific model and empirical evaluation, it may be necessary to either make it 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it 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 be a way to access this model for reproducing the results or a 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 to 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., to 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 to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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.cc/public/guides/CodeSubmissionPolicy ) 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 a subset of experiments are reproducible, they 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 as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all the training and test details in §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 a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in 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 not report error bars because the training behavior is consistent across different runs, and the goal of the experiments is to corroborate our main theoretical results, rather than achieving better performance on 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, or 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, or 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in 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 to 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 or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

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 do not see immediate societal impact of our paper.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or 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 (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If 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 ML).

11. 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 or models that have a 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

Question: 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?

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 use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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/model 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 is used.

At submission time, remember to 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 nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in 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 an 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 not 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors mention in the paper that Section C.1 contains a discussion of limitations, but this section is either missing or not clearly labeled.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: 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.

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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nQ1: 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?\n\nLIMITATIONS:\nA limitation of this study is that the model and data settings analyzed do not fully align with real-world transformers or language models.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSimilar 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?)\n\nLIMITATIONS:\nThe authors don't analyze their limitations", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-21"}