Title: Transformers on Markov data: Constant depth suffices
Venue: NeurIPS 2024
Accepted: yes
Keywords: Depth

k-gram

Transformers

Representation
Abstract: Attention-based transformers have been remarkably successful at modeling generative processes across various domains and modalities. In this paper, we study the behavior of transformers on data drawn from $k^{\text{th}}$-order Markov processes, where the conditional distribution of the next symbol in a sequence depends on the previous $k$ symbols observed. We observe a surprising phenomenon empirically which contradicts previous findings: when trained for sufficiently long, a transformer with a fixed depth and $1$ head per layer is able to achieve low test loss on sequences drawn from $k^{\text{th}}$-order Markov sources, even as $k$ grows. Furthermore, this low test loss is achieved by the transformer’s ability to represent and learn the in-context conditional empirical distribution. On the theoretical side, we prove that a transformer with $O(\log_2(k))$ layers can represent the in-context conditional empirical distribution by composing induction heads to track the previous $k$ symbols in the sequence. Surprisingly, with the addition of layer normalization, we show that a transformer with a constant number of layers can represent the in-context conditional empirical distribution, concurring with our empirical observations. This result provides more insight into the benefit of soft-attention and non-linearities in the transformer architecture.

Transformers on Markov data: Constant depth suffices

Abstract

Attention-based transformers have been remarkably successful at modeling generative processes across various domains and modalities. In this paper, we study the behavior of transformers on data drawn from k th -order Markov processes, where the conditional distribution of the next symbol in a sequence depends on the previous k symbols observed. We observe a surprising phenomenon empirically which contradicts previous findings: when trained for sufficiently long, a transformer with a fixed depth and 1 head per layer is able to achieve low test loss on sequences drawn from k th -order Markov sources, even as k grows. Furthermore, this low test loss is achieved by the transformer's ability to represent and learn the in-context conditional empirical distribution. On the theoretical side, our main result is that a transformer with a single head and three layers can represent the in-context conditional empirical distribution for k th -order Markov sources, concurring with our empirical observations. Along the way, we prove that attention-only transformers with O p log 2 p k qq layers can represent the in-context conditional empirical distribution by composing induction heads to track the previous k symbols in the sequence. These results provide more insight into our current understanding of the mechanisms by which transformers learn to capture context, by understanding their behavior on Markov sources. Code is available at: https://github.com/Bond1995/Constant-depth-Transformers .

1 Introduction

Attention-based transformers have revolutionized the field of natural language processing (NLP) [1, 2] and beyond [3, 4], achieving significant performance gains across tasks like machine translation, text generation, and sentiment analysis. A key factor in their success is their ability to model sequences far more efficiently, and the ability to learn in-context [5, 6].

To understand this capability, a canonical approach is to sample the input from a k th -order Markov process , where the next symbol's conditional distribution depends only on the previous k symbols. Recent studies [7, 6, 8] have investigated the ability of transformers to learn Markov processes and establish that learning happens in phases. The transformer eventually learns to represent the conditional k -gram model, which is the in-context MLE of the Markov process.

The results in [6, 8] seem to suggest that for low depth transformers to learn Markov processes of order k , it is essential that the number of heads scale linearly in k . At first glance, this is a bit concerning - real world data generating processes often contain long-range dependencies. How is it that transformers succeed at capturing these kinds of long-range dependencies, while at the same time requiring so many heads to be able to capture the necessary context for k th -order Markov sources?

¨ ¨ ¨

2

0

1

1

X

n

`

1

1

0

1

0

To understand the nature of this phenomenon, we train low-depth transformers on k th -order Markov sources. These experiments result in two surprising empirical phenomena that seem to contradict previous findings: when trained for sufficiently long, p i q a 2 -layer, 1 -head transformer can learn k th -order Markov processes for k as large as 4 , p ii q a 3 -layer, 1 -head transformer is able to achieve low test loss on sequences drawn from k th -order Markov sources, even as k grows to be as large as 8 (Fig. 3). In both cases, the values of k for which the models appear to learn k th -order Markov sources are much higher than those predicted in prior experiments [6, 8]. This discrepancy shows that our understanding of the mechanisms used by transformers to learn k th -order Markov processes is not complete and raises a broader question:

What is the interplay between depth, number of heads and non-linearity in learning k th -order Markov processes?

In this paper, we approach this question from the point of view of representation power, and provide some partial explanations toward the phenomena illustrated previously.

Our main contributions are as follows:

We show, rather surprisingly, that the standard transformer architecture with 3 layers and 1 head per layer is capable of representing the conditional k -gram model (Definition 1), and thereby learn k th -order Markov models in-context.

Along the way to building up to this result, we consider the simpler family of attentiononly transformers and show that they can represent the conditional k -gram model with r log 2 p k ` 1 q s layers.

Under a natural assumption on the nature of the attention patterns learnt by the transformer, we then argue that for k ě 3 attention-only transformers need at least r 1 ` log 2 p k ´ 2 q s layers to represent a ' k th -order induction head' (Definition 2). Empirically, transformers are observed to learn k th -order induction heads whenever they achieve small test error [6].

The last result is a consequence of a more general tradeoff between the number of layers, L , and heads per layer, H , an attention-only transformer requires to represent a k th -order induction head, under a natural assumption on the learnt attention patterns. In conjunction, these results also reveal the role of non-linearities (aside from the softmax in the attention) in the transformer architecture. In particular, it appears that layer normalization plays a critical role in the ability of constant-depth transformers to learn the conditional k -gram model. Together with the experimental results mentioned previously, these results paint a more comprehensive picture about the representation landscape of transformers in the context of k th -order Markov processes.

Notation. Scalars are denoted by italic lower case letters like x, y and Euclidean vectors and matrices in bold x , y , M , etc. The notation 0 p ˆ q (resp. 1 p ˆ q ) refers to the all-zero (resp. allone) matrix. When it is clear from the context, we omit the dimensions of a matrix. Define r S s fi t 1 , 2 , . . . , S u for S P N . I p¨q denotes the indicator function and Unif p S q denotes the uniform distribution over a set S .

1.1 Related work

There is a large body of active research focused on studying different aspects of transformer models [9, 10, 11, 12]. Our work closely relates to the aspects of understanding the representation power of

transformers, and in-context learning. [13, 14, 15] study the representation capabilities of transformers and show properties such as universal approximation and Turing-completeness. Viewing transformers as sequence to sequence models, [16, 17] study their ability to model formal languages and automata. Along more related lines to our work, [18, 19] present logarithmic depth transformer constructions for representing a k -hop generalization of the notion of an induction head [20]. On the other hand the theoretical and mechanistic understanding of in-context learning [21] has received much attention lately [22, 23, 24, 25], focusing on different operating regimes and phases of learning. There are a few recent papers which study the behavior of transformers when trained on data generated from Markov processes, and generalizations thereof [5, 26]. In particular, [7, 8] study the optimization landscape of gradient descent in learning generalizations of Markov processes, and [6] present a study of how transformers learn to represent in-context k -gram models, focusing on different phases of learning.

2 Preliminaries

We provide the necessary background for Markov processes, the conditional k -gram model, and the transformer architecture.

2.1 Markov processes

Markov processes are one of the widely used models in sequence modeling [27]. The characterizing property of these processes is that at any time step, the future evolution is only influenced by the most recent states. More formally, a sequence p X n q n ě 1 is a k th -order Markov process on a finite state space r S s with the transition kernel P , if surely,

This property allows us to capture the conditional distribution at any position using only its previous k symbols. This motivates the notion of a conditional k -gram, its empirical counterpart, defined for any sequence p x 1 , . . . , x n q .

Definition 1 (Conditional k -gram model) . Given a sequence p x 1 , ¨ ¨ ¨ , x n q of length n in r S s n , the conditional k -gram model x Pr k p¨ | x 1 , ¨ ¨ ¨ , x n q corresponds to the in-context estimate of the distribution over symbols conditioned on the last k symbols, i.e. for x P r S s ,

which is defined only so long as the denominator is non-zero. This structure is illustrated in Figure 2a. It is well known that the conditional k -gram in Eq. (1) with Laplace smoothing corresponds to the Bayes optimal estimate of the next symbol probability, when the data is drawn from fixed Markov process sampled from a prior distribution [27].

In our experiments, we will consider k th -order Markov kernels sampled from a Dirichlet prior with parameter 1 . Namely, the transition P p¨| X 1 ' i 1 , ¨ ¨ ¨ , X k ' i k q is sampled independently and uniformly on the S -dimensional simplex ∆ S 1 , for each tuple p i 1 , ¨ ¨ ¨ , i k q .

2.2 Transformer architecture

In this paper, we will consider variants of the standard transformer architecture in Figure 2b introduced in [1], with the goal to understand the role of depth and the non-linearities in the architecture. The simplest variant removes all the layer normalization and the (non-linear) feedforward layer, and is referred to as an attention-only transformer. The L -layer 1 -head attention-only transformer with relative position encodings, operating on a sequence of length T is defined in Architecture 1.

The attention scores in layer ℓ , t att p ℓ q n,i : i ď n u , are computed as Softmax `␣@ W p ℓ q K p x p ℓ q j ` p p ℓ q ,K n ´ j q , W p ℓ q Q x p ℓ q j D : j P r n s (˘ . The superscript p ℓ q indicates the layer index, and the matrices W p ℓ q K , W p ℓ q Q , W p ℓ q V P R d ˆ d capture the key, query and value matrices in layer ℓ . Note that the attention-only transformer may include a feedforward layer with linear activations, i.e. a linear

Input:

0

1

0

1

0

3

2

3

¨ ¨ ¨

0

1

match final

k

symbols

Uniform

distribution

(a) Conditional k -gram model. The conditional k -gram is the in-context estimate of the Markov process and is realized in two steps. The first step is to find the locations in the sequence (marked red) which match the final k symbols (functionally, a k th -order induction head). The conditional k -gram model returns the uniform distribution over the next symbol at these locations (marked blue).

Transformer layer

Layer norm

x

p

ℓ

q

Multi-head

Attention

+ POS

r

x

p

ℓ

q

Layer norm

FFN

x

p

ℓ

`

1

q

Layer

1

.

.

.

Layer

ℓ

.

.

.

Layer

L

x

p

1

q

Linear +

Prediction

Pr

x

1

,

¨ ¨ ¨

, x

T

q

Emb

Emb

Emb

¨ ¨ ¨

x

1

x

2

x

T

transformation. For representation purposes, this linear transformation can be combined with the projection matrix in the attention layer, allowing the feedforward layer to be omitted from the model. In the attention layer, we consider relative position encodings (the terms labeled in blue), which translates the key and value vectors depending on the relative position of the embedded symbol.

Architecture 1: Attention-only transformer.

The extension to H heads is straightforward, where in each transformer layer there are H attention layers in parallel, resulting in y p ℓ, 1 q n , ¨ ¨ ¨ , y p ℓ,H q n P R d for each n . These vectors are concatenated and passed through a linear transformation W p ℓ q O : R dH Ñ R d which is the output of the attention layer. Finally, the output of the model after L layers is passed through a linear layer, which projects the d -dimensional embeddings back into R S and the resulting vector is passed through a non-linearity f , usually a softmax, to result in the model's prediction of the next symbol probabilities. The theoretical results in this paper will choose f ' ReLU p¨q .

3 Understanding the empirical behavior of transformers

The motivation for the present work comes from a series of experimental results, which challenge our current understanding of transformers in the context of learning Markov processes. Several works in the literature [7, 8, 6] have studied the ability of transformer models to learn k th -order Markov processes. The experimental results present in the literature suggest that in order for a 2 layer transformer model to be able to learn a randomly sampled Markov process of order k , it is crucial for the number of heads in the first attention layer to scale linearly with the order, k . In particular, the authors of [6] claim that in their experiments, 'Single attention headed models could not achieve better performance than bigram (models)' in learning random k th -order Markov processes in-context. Similarly, the

0

5000

10000

15000

0.1

0.2

0.3

0.4

Test loss gap from the optimal

k

'

2

k

'

4

0

5000

10000

15000

20000

25000

0.0

0.1

0.2

0.3

0.4

0.5

Test loss gap from the optimal

k

'

2

k

'

4

k

'

8

Iteration

Iteration authors of [8] study a generalization of learning k th -order Markov processes to learning causal processes on degree k graphs. The theory and experiments pertain to 2 -layer k head transformers.

In Figure 3, we train 2 and 3 -layer transformers with a single head on data drawn from random Markov processes of various orders drawn from a Dirichlet prior. With 2 layers and a single head, we see that the model is able to learn even order4 Markov processes, and go beyond the simple order1 processes which were projected to be the limit of its ability to learn. Likewise, with 3 layers, transformers are able to go much further and learn order8 Markov processes, which was the largest value of k we evaluated on. These results contrast with our current understanding of how induction heads are realized in the parameter space [6, 8] - existing constructions which realize these attention patterns require k heads when the number of layers is 2 , and it's unclear how to implement them with fewer heads. At a high level, each of the k heads play a critical role - where, loosely speaking, the i th -head looks back i positions in the sequence.

Building up to our main results, in the sequel, we study the simpler case of attention-only transformers where the feedforward layers and layer normalization are removed.

4 Warming up: Attention-only transformers

The study of attention-only transformers trained on Markov processes has garnered some attention in the prior literature. Notably, the authors of [6] study 2 -layer 1 -head attention-only transformers trained on data drawn from 1 st -order Markov processes whose parameters are drawn from a Dirichlet prior. The model is observed to learn a very specific behavior, known as an 'induction head' [20], which in this setting is able to represent the conditional 1 -gram (Eq. (1)).

The induction head mechanism is composed of two layers where the first layer learns the attention pattern att p 1 q n,i ' I p i ' n ´ 1 q , thereby allowing the model to capture information about the symbol at position n ´ 1 in the embedding vector at time n . In the second layer, the attention layer picks out those indices n where x n ´ 1 ' x T , the final symbol in the sequence. At these positions, since x n ´ 1 ' x T , one would expect that the next symbol x n is a good predictor of x T ` 1 , and the model uses this information to predict the next symbol x T ` 1 according to its conditional empirical estimate, x Pr 1 p x T ` 1 | x 1 , ¨ ¨ ¨ , x T q , i.e. the conditional 1 -gram model.

Theorem 1. The conditional 1 -gram model can be represented by a 2 -layer and 1 -head attention-only transformer with embedding dimension d ' 3 S ` 2 .

Although a version of this result is proved in [6], we include a proof in Appendix A for completeness.

Remark 1. In Theorem 1 and other results to follow, we de-emphasize the role of the bit-precision to which the transformer is implemented. That said, note that when the constructions in Theorems 1 to 3 are implemented to O p log p T qq bits of precision, the representation results are realized up to an additive O p 1 { T q error.

The ideas in Theorem 1 readily extend to representing the conditional k -gram model, by instead using k heads in the first layer. The j th head learns the attention pattern att p 1 q n,i ' I p i ' n ´ j q and

concatenating the outputs of the heads, the model learns to aggregate information about x n , ¨ ¨ ¨ , x n ´ k in the embedding vector at time n . The second layer realizes what is best described as a ' k th -order ' induction head, where the model learns to pick out those positions n where for every j P r k s , x n ´ j ' x T ´ j ` 1 , i.e. the history of length k at those positions match the final k symbols in the input sequence ( see Figure 4). This mechanism is also referred to as a long-prefix induction head [28].

Definition 2 (Higher-order induction head) . A 1 -head attention layer is said to realize a k th -order induction head if on any sequence p x 1 , ¨ ¨ ¨ , x T q P r S s T , for any fixed n ď T , as a function of the input sequence, att n,T is maximized if and only if x n ´ j ' x T ´ j ` 1 for every j P r k s .

k th -order induction heads generalize the concept of an induction head [20], and keep track of the positions i ď n where there is a perfect occurrence of the final k symbols in the sequence. Such attention patterns are immediately useful in representing the conditional k -gram - increasing the temperature within the softmax of this attention layer results in an attention pattern which converges to the uniform distribution over those positions where the final k symbols x T ´ k ` 1 , ¨ ¨ ¨ , x T are seen previously in the sequence. Loosely, this allows the model to 'condition' on the last k symbols in the sequence. With k heads, the model can aggregate information from the previous k positions and implement a k th -order induction head, which leads to the following result. A full proof is discussed in Appendix A.1.

Input:

n

'

0

1

0

1

0

3

2

3

1

2

3

6

4

5

7

8

9

10

10

1

2

3

4

5

6

7

8

9

0

0

.

5

1

n

att

T,n

0

1

Theorem 2. The conditional k -gram model can be represented by an attention-only transformer with 2 layers, k heads and embedding dimension d ' p k ` 2 q S ` k ` 1 .

While this result is positive, it suggests that a 2 -layer transformer requires approximately k times as many parameters to be able to represent the conditional k -gram model. The first result we prove is that increasing the depth of the model is exponentially more beneficial, in that a transformer with O p log p k qq depth can estimate in-context k -grams.

Theorem 3. The conditional k -gram model can be represented by an attention-only transformer with relative position encodings, with L ' r log 2 p k ` 1 q s layers and 1 head per layer. The embedding dimension is ď 2 k p S ` 1 q ` S .

With 2 layers and k heads, the transformer aggregates information about each of the previous k positions one step at a time through the k heads. However, with Ω p log p k qq layers, the same task can be done far more efficiently. In the first attention layer, the model aggregates information about the current and previous position. Namely, using the relative position embeddings, x p 2 q n is chosen as a linear combination of x p 1 q n ' Emb p x n q and x p 1 q n ´ 1 ' Emb p x n ´ 1 q . This allows the embedding at position n to aggregate information about x n and x n ´ 1 . In the same vein, in the second attention layer, the model aggregates information from x p 2 q n and x p 2 q n ´ 2 in x p 3 q n ; the former has information about x n and x n ´ 1 , and the latter has information about x n ´ 2 and x n ´ 3 . This expands the 'window' of x i 's on which x n depends on to size 4 . In the ℓ th layer, the model aggregates information from x p ℓ q n and x p ℓ q n ´ 2 ℓ which allows x p ℓ ` 1 q n to effectively depend on the x i 's in a window of size 2 ℓ ` 1 starting at position n , namely x n , ¨ ¨ ¨ , x n ´ 2 ℓ ` 1 ` 1 . In the final layer, the embedding at position i , x p L q i for L ' r log 2 p k ` 1 q s depends on x n , x n ´ 1 , ¨ ¨ ¨ , x n ´ k . In the last layer, the model can realize the dot-product A W p L q K x p L q n , W p L q Q x p L q T E ' ř k j ' 1 I p x n ´ j ' x T ´ j ` 1 q by choosing the key and query vectors appropriately. By increasing the temperature in the attention softmax, the attention pattern realized is the uniform distribution on values of n such that x n ´ j ' x T ´ j ` 1 for every j P r k s , i.e., a k th -order induction head. The full proof of this result is provided in Appendix B.

Multi-head

Attention

+ POS

x

p

ℓ

q

Ă

xy

p

ℓ

q

FFN

Layer norm

x

p

ℓ

`

1

q

Layers

1

`

2

»

-

-

-

Emb

p

x

n

q

Emb

p

x

T

q

¨ ¨ ¨

u

n

}

u

n

}

2

¨ ¨ ¨

u

T

}

u

T

}

2

v

n

}

v

n

}

2

v

T

}

v

T

}

2

fi

ffi

ffi

fl

Layer

3

r ¨ ¨ ¨

Emb

p

x

i

q

¨ ¨ ¨

Emb

p

x

T

q s

att

T,i

9

exp

´

κ

x

v

i

,

u

T

y

}

v

i

}

2

}

u

T

}

2

¯

(realizes a

k

th

-order induction head)

Emb

x

1

,

¨ ¨ ¨

, x

T

While this is a promising step toward understanding the behavior transformers exhibit in Figure 3, showing that depth plays an important role in their ability to represent conditional k -gram models, the picture is still not complete. The experimental results in Section 3 do not preclude the possibility that a transformer might not even require logarithmic depth to be able to learn k th -order Markov processes approximately. In the next section, we will study constant-depth transformers and establish a rather surprising positive result about the representation power of this class of models in capturing conditional k -grams.

5 Understanding the role of non-linearity: Constant-depth constructions

In the previous section, we saw how the transformer uses the power of depth to learn conditional k -grams far more efficiently. In particular, every additional attention layer effectively doubles the window of positions i ' n ´ 1 , n ´ 2 , ¨ ¨ ¨ which the model has access to information about at the current time n . By composing L ' Ω p log p k qq attention layers, the model is able to collect enough information within the output embedding x p L ` 1 q n to be able to realize a k th -order induction head in the next layer. In this section, we prove that adding non-linearity to the architecture, in the form of layer normalization, can significantly change the mechanism in which the transformer realizes this k th -order head. In particular, there are constant depth architectures which allow a k th -order induction head to be realized, surpassing the logarithmic depth attention-only constructions.

Modification to the standard transformer architecture. To simplify the proof of our main result, we will consider a subtle modification to the standard transformer architecture, which is presented in Architecture 2 and Figure 5a. We will remove the first layer norm prior to the multi-head attention and move the second layer norm to after the feed-forward network. It is important to note that Theorem 4 holds even for the architecture presented in Figure 2b, which is the architecture we evaluate empirically. The modification we present in Figure 5a allows the construction to be simpler and makes it much easier to convey the key intuition. The main difference compared to the attentiononly design presented in Architecture 1 is the addition of layer normalization and a feedforward layer in the for-loop over n P r T s for each transformer layer ℓ . The differences between Architectures 2 and 1 are emphasized in blue.

Theorem 4. Conditional k -grams can be represented by a transformer with 3 layers, 1 head per layer, relative position encodings and layer normalization. The embedding dimension is O p S q .

Remark 2. Although the proof stated does not bound the approximation error arising from a finite bound on the bit precision of the transformer, in theory, it should suffice to have Ω p log p T q ` k q bits per parameter for the statement of Theorem 4 to go through with an O p 1 { T q additive approximation error. The main point is that none of the weights of the model exceed exp p k q and with log p T q additional bits per parameter, the approximation error scales as O p 1 { T q .

Architecture 2: Modified transformer architecture. The computations above are carried out for each n P r T s in each layer ℓ P r L s . In the layer normalization step (LN), the feature mean µ is defined as, E i ' Unif pr d sq '@ e d i y p ℓ q n D‰ and the feature variance σ 2 ' E i ' Unif pr d sq '@ e d i , y p ℓ q n D 2 ‰ ´ µ 2 .

5.1 Proof sketch

In the attention-only transformer with 2 layers and k heads, the model is able to keep track of where the final k symbols in the sequence appeared previously (i.e., a k th -order induction head) by, loosely, using each head to keep track of the occurrences of one of the final k symbols. On the other hand, with the benefit of more depth, with L ' Ω p log p k qq layers, the model is able to collect enough information within the output embedding x p L ` 1 q n to be able to realize the same behavior. However, neither of these constructions scale down to the case when the depth and number of heads of the transformer are both constants independent of k . We provide a brief sketch of the construction below.

Recall that a k th -order induction head keeps track of the indices i such that @ j P r k s , x i ´ j ' x n ´ j ` 1 . Defining z i fi ř k j ' 1 2 j e x i ´ j ` 1 , notice that the condition t@ j P r k s , x i ´ j ' x n ´ j ` 1 u can equivalently be captured by writing t z i ´ 1 ' z n u . This true because of the fact that the binary representation of any integer is unique. Furthermore, these vectors, up to scaling, can be realized by softmax attention (namely, att n,n ´ i 9 2 i for 1 ď i ď k ).

With this step, finding occurrences of the last k symbols in the input sequence boils down to realizing an attention pattern in the second layer, att p 2 q n,i , which is maximized whenever z i ´ 1 ' z n . While dot-product attention naively encourages those values of i for which z i ´ 1 and z n are 'similar' to each other, a qualitative statement is lacking. In general, it will turn out to that a different measure of similarity is necessary within the softmax to be able to encourage those values of i for which these vectors match. This is where the role of layer-normalization comes in.

Instead of the usual dot-product, suppose the attention mechanism in the second layer was,

where κ is the temperature parameter. Then, as the temperature κ grows, the attention pattern essentially focuses on those values of i for which z i {} z i ´ 1 } 2 ' z n {} z n } 2 . With this attention pattern, we are thus very close to the statement we wanted to check, ( z i ´ 1 ? ' z n ). As it turns out, for the special structure in the z i 's considered (dyadic sums of one-hot vectors), we may write down,

A quantifiable equivalence is provided in Lemma 1.

Realizing L 2 -norm attention (eq. (1) ). Observe the equivalence,

Taking a softmax on both sides, notice that the RHS (up to an additive constant) is the L 2 -norm based attention, while the LHS is the usual dot-product attention between z i ´ 1 {} z i ´ 1 } 2 and z n {} z n } 2 . Thus on unit-normalized vectors, L 2 -norm attention and dot product attention are but the same.

While the first layer of the transformer computes the z i 's by a weighted summation, layer normalization fills in the last missing piece of the puzzle which is to normalize them to unit norm. This is a consequence of defining the embedding vectors appropriately, as we discuss more in Appendix C.1.

From this step, realizing the actual conditional k -gram model follows readily. In particular, as the temperature κ in the attention grows, the attention pattern zooms in on indices i P I n fi t k ` 1 ď i ď n : @ j P r k s , x i ´ j ' x n ´ j ` 1 u in the last layer. The value vectors at this step are the one-hot encoding of x i ; putting everything together, the logits realized by the transformer are,

which is the conditional k -gram model (eq. (1)).

While the transformer construction described above only requires two layers, the actual construction we propose differs slightly and has an additional layer. The first two layers of the transformer respectively compute z i and z i ´ 1 which are added to the embedding vector at time i . This is important because we need to test whether z i ´ 1 ? ' z n and not whether z i ? ' z n or z i ´ 1 ? ' z n .

Summary. The construction can be summarized as follows: the first layer computes z n ' ř k j ' 1 2 j ´ 1 ¨ e x n ´ j by choosing appropriate value vectors and relative position embeddings to realize the attention pattern attn n,n ´ i 9 2 i I p 1 ď i ď k q . The layernorm that follows subsequently can be replaced by RMSnorm, by a simple trick which we discuss in Appendix C.1, resulting in z n {} z n } 2 to be appended to the embedding at time n . Using a very similar construction, layer 2 computes z n ´ 1 {} z n ´ 1 } 2 , which is added to the embedding at time n . Finally, in the last layer, the dot-product x z i ´ 1 {} z i ´ 1 } 2 , z n {} z n } 2 y defines the attention score, and as the temperature κ grows, the pattern converges to Unif p I n q . Choosing the value vectors in this layer appropriately gives eq. (3).

6 Lower bounds on transformer size

In this section, we study the limits of how shallow a transformer can be made while still capturing conditional k -grams. The first result we establish in this vein is a lower bound against 1 -layer transformers showing that their expressive power is too limited unless the embedding dimension or number of heads scale near-linearly in T .

Theorem 5. Consider any 1 -layer transformer with layer normalization and feedforward layers, where all the coordinates of the embedding vectors and unnormalized attention scores are computed with p bits of precision. If the transformer is able to compute the conditional 3 -gram on inputs drawn from t 0 , 1 , 2 u T to within an additive error of 1 { 3 T , then 2 pH ` dp ` 2 ě T { 3 .

Choosing the bit precision to be p ' O p log p T qq , this implies that for transformers with 1 layer, the sum of the number of heads and the embedding dimension must be at least Ω p T { log p T qq , in order to represent conditional 3 -grams to within an additive error of 1 { 3 T .

6.1 Conditional lower bounds on attention-only transformers

While the previous section shows that 1 -layer transformers have fairly limited representation power, it is not immediately clear how whether any of these issues are present with transformers with more layers. Indeed, as we discussed in Section 4, an attention-only transformer with O p log 2 p k qq layers and 1 head per layer can represent conditional k -grams on its input sequences. With the addition of non-linearities, Theorem 4 shows that the model can represent conditional k -grams using just a constant number of layers. In this section, we try to understand the gap between these two results and prove conditional lower bounds on the size of attention-only transformers which do not have non-linearities arising from layer normalization.

We prove conditional lower bounds under some natural assumptions on the nature of the attention patterns learnt by the transformer. To motivate these assumptions, consider the experiment in Figure 6, where we train an attention-only transformer with 2 layers and 1 head, on order1 Markov processes. At test-time, we plot the attention patterns learnt in the first layer of the model on test sequences. Notice that the attention pattern learnt by the model at layer 1 is largely independent of the input sequences themselves and only depends on the position.

Assumption 1. In an L -layer attention-only transformer with H heads per layer, assume that layers ℓ ' 1 , 2 , ¨ ¨ ¨ , L ´ 1 and heads h P r H s realize an attention pattern where att p ℓ,h q n,i only depends on the positions n and i and on ℓ and h , but not on the input sequence x 1 , ¨ ¨ ¨ , x T .

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0.0

0.2

0.4

0.6

0.8

1.0

Index

i

Index

n

0

5

10

15

20

25

30

0

5

10

15

20

25

30

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Index

i

Index

n

0

2

4

6

8

10

0.0

0.2

0.4

0.6

0.8

1.0

Sequence index

p

i

q

Attention weight

Rather than proving the size lower bound depending on the transformers ability to represent the conditional k -gram itself, we consider a simplification and assume that the goal of the model is to represent a k th -order induction head (Definition 2) in the last layer. Although learning a k th -order induction head is not strictly necessary for the transformer to be able to represent conditional k -grams, note that every construction we have considered so far (cf. Theorems 1 to 4) go through this mechanism to realize the conditional k -gram model. Likewise, for other related problems, such as the causal learning task in [8], the causal structure is captured by an extension of the k th -order induction head to general causal graphs. Our main lower bound is the following result.

Theorem 6. Consider an L -layer transformer with h ℓ heads in layer L . Assuming the transformer satisfies Assumption 1, if ś L ´ 1 ℓ ' 1 p H ℓ ` 1 q ď k ´ 2 , the attention pattern in layer L cannot represent a k th -order induction head.

While this lower bound is not unconditional, meaning that it does not directly imply that the transformer cannot represent conditional k -grams, it is important to understand the interpretation of this result: attention-only transformers which somehow break through this barrier need to use a significantly different mechanism to realize the conditional k -gram model.

Theorem 6 implies that under Assumption 1, a 2 -layer attention-only transformer with 1 head cannot realize a k th -order induction head for any k ě 4 . Likewise, under the same assumption, a 3 -layer attention-only transformer with 1 head cannot realize a k th -order induction head for any k ě 6 . These results give more weight to the experiment in Figure 3 where we observe that a 2 -layer transformer learns a k th -order Markov process for k ' 4 and a 3 -layer transformer learns a k th -order Markov process for k ' 8 , and show that non-linearities in the architecture allow the transformer to break past the size barriers in Theorem 4.

7 Conclusion

We observe empirically that 2 and 3 layer transformers are able to learn k th -order Markov chains for much higher values of k than previously anticipated. We show there are O p log p k qq -layer constructions of attention-only transformers which are able to learn the conditional k -gram model, which is the in-context MLE of the Markov model. With non-linearities in the model, we show that a 3 -layer 1 -head transformer is capable of representing the same. We show that 1 -layer transformers cannot represent conditional k -grams for any k ě 3 unless the number of heads or embedding dimension scale almost linearly in T . We also prove a conditional lower bound on the depth and number of heads of attention-only transformers to represent k th -order induction heads, under an assumption on the realized attention patterns.

Acknowledgments and Disclosure of Funding

The work was partially supported by NSF Grant CCF-2211209 and Swiss NSF Grant 200364.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998-6008, 2017.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners, 2021.

Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024.

Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.04161, 2024.

Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent, 2024.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11090, 2021.

Angeliki Giannou, Shashank Rajput, Jy-Yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In Proceedings of the 40th International Conference on Machine Learning, pages 11398-11442, 23-29 Jul 2023.

Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In Proceedings of the 40th International Conference on Machine Learning, 2023.

Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations, 2023.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations, 2020.

Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing-complete. Journal of Machine Learning Research, 22(75):1-35, 2021.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. In Advances in Neural Information Processing Systems, volume 35, pages 12071-12083, 2022.

Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata, 2023.

Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages, 2020.

Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth, 2024.

Clayton Sanford, Daniel J Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 36677-36707. Curran Associates, Inc., 2023.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.

Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning, 2024.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2023.

Jesse Hoogland, George Wang, Matthew Farrugia-Roberts, Liam Carroll, Susan Wei, and Daniel Murfet. The developmental landscape of in-context learning. arXiv preprint arXiv:2402.02364, 2024.

Nived Rajaraman, Jiantao Jiao, and Kannan Ramchandran. Toward a theory of tokenization in llms, 2024.

J. R. Norris. Markov Chains. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1997.

Nicholas Goldowsky-Dill, Chris MacLeod, Lucas Sato, and Aryaman Arora. Localizing model behavior with path patching. arXiv preprint arXiv:2304.05969, 2023.

Andrew Chi-Chih Yao. Some complexity questions related to distributive computing (preliminary report). In Proceedings of the eleventh annual ACM symposium on Theory of computing, pages 209-213, 1979.

Matteo Pagliardini. GPT-2 modular codebase implementation. https://github.com/epfml/llm-baselines, 2023.

Appendix

Table of Contents

Notation. The notation e d 1 i P R d 1 refers to the one-hot encoding of i in d 1 dimensions. In other words it is the i th standard basis vector in d 1 dimensions. The notation Blkdiag pt A 1 , A 1 , ¨ ¨ ¨ , A m uq refers to the block diagonal matrix with i th block as A i .

A Proof of Theorem 1

We will first prove Theorem 1. In the first layer, choose the embeddings as,

for a constant κ ą 0 to be chosen later and d ' 2 S ` 2 . The relative position encodings will essentially be supported on the first two coordinates, the middle S coordinates are a one-hot encoding of the symbol x n and the last 2 S coordinates are 0 . The relative position encodings in the first layer are chosen to be p p 1 q ,K n ´ i ' κ p´ 1 ` I p n ´ i ' 1 uqq e d 1 P R d and p p 1 q ,V n ´ i ' 0 P R d . Choose W p 1 q K

and W p 1 q Q to be e d 1 p e d 1 q T P R d ˆ d . With this choice,

As κ Ñ8 , the attention pattern (which takes the softmax over of these inner products over i P r n s ) computes,

for any n ą 1 . Choose the value matrix as,

And with this choice and the residual connection, we get,

which serves as the input to the 2 nd transformer layer.

Layer 2. In layer 2 , the relative position encodings p K, p 2 q n ´ i and p V, p 2 q n ´ i are all set as 0 . The key matrix picks out the e S x n block out of x p 2 q n and the query vector picks out the e S x i ´ 1 block out of x p 2 q i ´ 1 . In particular, these matrices are chosen so that,

Taking the inner product of these vectors, and taking κ Ñ 8 , observe that the attention pattern concentrates on the uniform distribution over all coordinates i such that x i ´ 1 ' x n . More formally, the attention pattern for any n ą 1 is,

assuming ř n i ' 2 I p x i ´ 1 ' x n q ą 0 . Having realized this attention pattern, may choose the value and subsequent linear layer appropriately. The value matrix simply picks out the e S x i block from x p 2 q i and places it into the last S coordinates of x p 3 q i , and the linear layer simply extracts this block and outputs it (after scaling down by a factor of κ ), realizing the logits,

if ř n i ' 2 I p x i ´ 1 ' x n q ą 0 . In particular, under the same condition,

assuming ř n i ' 2 I p x n ´ 1 ' x T q , which is the conditional 1 -gram model.

A.1 Extension to k -heads: Proof of Theorem 2

In the first layer, the embeddings are chosen to be,

With d ' p k ` 1 qp S ` 1 q ` S . The relative position encodings are chosen as p K, p 1 q i ' ' e k i 0 ‰ T for 1 ď i ď k and p K, p 1 q i ' 0 otherwise. Similarly, p V, p 1 q i ' 0 for every i . The h th head has key and query matrices,

With these choices, and letting κ Ñ8 , the h th layer computes the attention pattern,

Choose the corresponding value matrix as,

choosing the projection matrix appropriately, the output of the transformer after the first residual connection is,

Layer 2. In this layer, the relative position encodings p K, p 2 q n ´ i and p V, p 2 q n ´ i are all set as 0 . The key and query matrices are chosen as,

With this choices, we have that,

Taking κ Ñ8 , observe that the attention pattern concentrates on the uniform distribution over all coordinates i such that x i ´ j ' x n ´ j ` 1 for all j P r k s . More formally, if ř n i ' 2 I p x i ´ 1 ' x n q ą 0 , the attention pattern for any n ą 1 is,

The value matrix picks out e S x i from the embedding x p 2 q i (Equation (17)) and places it in the last S coordinates. The subsequent linear layer picks out the last S coordinates, resulting in the logits,

assuming that ř n i ' k ` 1 I p@ j P r k s , x i ´ j ' x n ´ j ` 1 q ą 0 . In particular,

assuming ř T n ' k ` 1 I p@ 1 ď j ď k, x n ´ j ' x T ´ j ` 1 q ą 0 , i.e., the conditional k -gram model.

B Proof of Theorem 3

Define k ‹ ' 2 r log 2 p k ` 1 q s by rounding k ` 1 up to the nearest power of 2 and ℓ ‹ ' log 2 p k ‹ q . In the setting of relative position encodings, given the sequence x 1 , ¨ ¨ ¨ , x n , while generating the output of the attention + feedforward layer for the symbol x n , the embeddings x i ' Emb p x n q ` p n ´ i are used for i P r n s . In other words, the position encoding vector is taken relative to the end of the sequence, rather than the start of the sequence. Consider the embedding of x as,

where e d 1 i P R S is the standard basis vector in d 1 dimensions. And the relative position encoding for the keys as,

And for the value vectors, p V i ' 0 for all i .

For the first layer and first head, we will describe the value, key and query matrices. Choose,

Then, observe that for i ě 1 ,

In particular, letting κ Ñ8 , the attention pattern is,

together with the residual connection, we get,

Layer ℓ ` 1 . By induction, assume that the output of the ℓ th transformer layer is of the form,

for some vector v n P R 2 ℓ S . We will show that with appropriately chosen key, query and value vectors in the p ℓ ` 1 q th layer, the output of this layer is,

We will consider the same relative position encodings and query matrix in this layer as in the first layer (Equations (24) and (25)). Consider a key matrix of the form,

With this choice, observe that for i ě 1 ,

and for i ' 0 ,

In particular, letting κ Ñ8 , the attention pattern is,

and for i ' 0 ,

Choose the value matrix as,

Choosing the value matrix as, we get,

Final last transformer layer ( ℓ ' ℓ ‹ ). The output of the second last transformer layer, indexed ℓ ‹ ´ 1 is,

which follows by plugging in the definition of k ‹ . Note that there exists a linear transformation L p ℓ ‹ q such that,

This can be further decomposed as,

And yet again there exists a linear transformation L p ℓ ‹ ´ 1 q which transforms this as,

By recursing this argument and composing all the linear transformations, up to a global permutation, we get that,

In the final layer, we will right multiply the key, query and value matrices by L ‹ ' ś ℓ ‹ ℓ ' 1 L p ℓ q . The effect can be interpreted as operating the original key, query and value matrices on the embedding vectors in Equation (33). In the final layer, we will set all the position encodings to be 0 and consider the key and query matrices,

Then,

Where we must be careful to note that the input x p ℓ ‹ q n contains copies of e x n , e x n ´ 1 , ¨ ¨ ¨ , e x n ´ k since k ‹ ě k ` 1 by definition.

Letting κ Ñ 8 , if there exists i such that ř k ´ 1 j ' 0 I p x n ´ j ' x i ´ j ´ 1 q ą 0 , for n ě k , the attention pattern is,

Finally, choose, we get,

Choosing the subsequent linear layer as,

Results in the output,

which is precisely the in-context conditional k -gram.

C Proof of Theorem 4

C.1 Modifying the definition of layer normalization

In every layer, we will perform a simple transformation which is to double the hidden dimension d and add a copy of ´ x p ℓ q n into the last d coordinates. This is possible by modifying the weights of the transformer appropriately as discussed below. A consequence of this transformation is that the feature mean of the x n 's is µ n ' 0 , and therefore the standard deviation σ n simply normalizes by the L 2 -norm of the features. In order to avoid having to explicitly state this transformation at each layer, we will simply redefine the layer norm LN to output v {} v } 2 for the input vector v , which is realized on the first d coordinates of the transformed embeddings.

This transformation can be realized automatically by redefining the initial embeddings Emb p x n q , and modifying the weights of the attention and feedforward subnetworks as follows: The input embeddings are changed to r Emb p x n q ´ Emb p x n qs T P R 2 d . The key and query matrices are chosen to be 0 on the last d coordinates in every layer; the value matrix for i ě 1 is transformed to Blkdiag `␣ W p ℓ q V , W p ℓ q V (˘ , and likewise changing the feedforward layer to the block diagonal matrices Blkdiag `␣ W p ℓ q 1 , W p ℓ q 1 (˘ and Blkdiag `␣ W p ℓ q 2 , W p ℓ q 2 (˘ . This transformation adds a copy of ´ x p ℓ q n into the last d coordinates of the corresponding embeddings.

C.2 Notation and supplementary lemmas

For each i P r T s , define,

Note that although v i ' u i ´ 1 , we make the distinction between the two to avoid any confusion in what is stored in the embedding vector at time i and at time i ´ 1 . Furthermore, define,

Lemma 1. If i P I n , z i ' z n ´ 1 . However, if i ě k ` 1 but i R I n , then, › › v i } v i } 2 ´ u n } u n } 2 › › 2 ě 3 ´ k .

Let j ‹ P t 0 , 1 , ¨ ¨ ¨ , k ´ 1 u denote the largest index j such that x n ´ j ‰ x i ´ j ´ 1 . Consider the coordinates a ' x n ´ j ‹ P r S s and b ' x i ´ j ‹ ´ 1 P r S s . Then,

If } v n } 2 ě } u i } 2 , then,

This uses the fact that u i and v n are coordinate-wise non-negative. On the other hand, if } v n } 2 ď } u i } 2 , using a similar analysis,

In either case, there is a coordinate ( a or b ) such that, u i {} u i } 2 and v n {} v n } 2 differ by at least 3 j ‹ ´ k . This implies the lower bound on the L 2 norm of the difference of the vectors.

C.3 Proof of Theorem 4

Choose the input embeddings as,

In the first two layers we will use the same relative position embeddings, in particular,

and the value embeddings,

Then, observe that,

Letting κ Ñ8 , this results in the attention pattern,

The output of the attention layer (with the residual connection) is,

Choose the value matrix as,

In the final layer, we will drop all position-related information and choose p p 3 q ,K i ' p p 3 q ,V i ' 0 for all i .

Layer 1. Consider the key and query matrices,

In the feedforward layer to follow, we will choose,

Which simply extracts u n from r x p 1 q n . With the subsequent layer norm and residual connection, the output of the first layer is,

Layer 2. In this layer, the relative position encodings and query matrix are the same as in layer 1 but the key matrix is chosen as,

With this choice, observe that,

As before, since κ Ñ8 , this results in the attention pattern,

which is similar, but subtly different from the attention pattern in the first layer (Equation (53)). The first layer focuses on indices n ´ i such that 0 ď i ď k ´ 1 , while this layer focuses on 1 ď i ď k . Choosing the value and projection matrices as,

The output of the attention layer (with the first residual connection) is,

It is a short calculation to see that Z n ' 3 k ` 1 { 5 if n ě k ` 1 and otherwise, Z n ď 3 k { 5 . This will be useful later, since the value of Z n can be used to determine whether n ě k ` 1 or n ď k which will allow the the next layer to avoid calculating the attention at i ď k , where the evaluation x n ' x i ´ 1 , ¨ ¨ ¨ , x n ´ k ` 1 ' x i ´ k is not well defined. In the subsequent FFN layer, we will choose,

Which extracts v n from the embedding r x p 2 q n . With the layer norm and adding the final residual connection, the output of this layer is,

Layer 3. In this layer, all the relative position encodings are set as 0 and instead,

With these choices,

The resulting attention scores are,

Recall that I n ' t k ` 1 ď i ď n : @ j P r k s , x n ´ j ` 1 ' x i ´ j u . Then for any i P I n , v i ' u n , and by Lemma 1, for any i ě k ` 1 but not in I n ,

Note that this gap is small but non-zero. Furthermore, recall that Z i ' 3 k ` 1 { 5 if i ě k and otherwise Z i ď 3 k { 5 . Thus the attention prefers values of i such that v i ' u n and such that i ě k ` 1 . In particular, as κ Ñ8 , the resulting attention pattern is,

Choosing,

We get that,

The feedforward layer is chosen to have W p 3 q 1 ' W p 3 q 2 ' 0 , and the overall output of the final transformer layer is therefore just r x p 3 q n . In the output linear layer, choose,

which results in,

In particular,

which is the conditional k -gram.

D Representation lower bounds for 1 -layer transformers: Proof of Theorem 5

We prove this lower bound by a reduction to communication complexity, and specifically to the set disjointness problem.

Suppose Alice and Bob are given strings a , b P t 0 , 1 u n which are indicator vectors of sets A and B . Their goal is to jointly compute DIS p a , b q ' I pD i : a i ' b i ' 1 q , which indicates whether A and B intersect or not. Alice and Bob may send a single bit message to the other party over a sequence of communication rounds. The following seminal result by [29] asserts a lower bound on amount of communication required between Alice and Bob to carry out this task.

Theorem 7 ([29]) . Any deterministic protocol for computing DIS p a , b q requires at least n rounds of communication.

We show that a 1 -layer transformer with sufficiently small embedding dimension / number of heads can be used to simulate a two-way communication protocol between Alice and Bob to solve DIS p a , b q in a way which contradicts Yao's lower bound in Theorem 7.

With m ' T { 3 ´ 1 , suppose Alice and Bob have length m bit strings a , b P t 0 , 1 u m . The transformer's input will be a sequence of the form,

of length 3 m ` 2 ' T ´ 1 . The input basically contains a repeating motif, composed of the symbol 2 followed by one of Alice's bits, and then one of Bob's bits. The last 2 symbols are 2 and 1 . We will consider the empirical conditional 3 -gram probability the transformer associates with the symbol x T ' 2 . Noting that x T ´ 1 ' 1 and x T ´ 2 ' 1 , the conditional 3 -gram is computed to be,

Note that if x i ´ 2 ' 2 , then i must be of the form 3 j for j ' 1 , ¨ ¨ ¨ , n , and we may rewrite the sum as,

Now, let us use the transformer to construct a deterministic communication protocol between Alice and Bob. Alice is given p x 2 , x 5 , ¨ ¨ ¨ , x 3 m ´ 1 q ' p a 1 , a 2 , ¨ ¨ ¨ , a m q and Bob is given p x 3 , x 6 , ¨ ¨ ¨ , x 3 m q ' p b 1 , b 2 , ¨ ¨ ¨ , b m q .

In the first round, Alice computes the normalization in the softmax of the attention which comes from the set of inputs she holds. For simplifying notation define,

In particular, for each head h P r H s , she computes,

Assuming that the transformer uses p bits of precision, Alice communicates Z p h q Alice for each h , which corresponds to pH bits of communication. With this information, Bob completes the rest of the normalization term (again up to p bits of precision) and computes,

which is the overall normalization term in the softmax. This is communicated back to Alice, using another pH bits of communication. Next using this information, Alice computes the output of the

attention layer, taking the convex combination corresponding to the inputs she knows. In particular, for each h P r H s she computes,

across all the heads. Rather than transmitting everything, she concatenates the outputs of the heads, and multiplies them by the value and projection matrices to result in the output y Alice which is d -dimensional. This is sent to Bob using dp bits of communication. Subsequently, Bob computes the terms in the attention corresponding to the inputs he knows as well as the public inputs (all the 2 's at positions 3 j ´ 2 as well as the last two symbols). In particular,

These are yet again concatenated across all the heads and multiplied by the value and projection matrices to result in the output y Bob which is added to y Alice to result in y . Bob passes y through the residual connection, layer norm, and feedforward layers, and subsequently through the linear layer and softmax of the model to result in the output of the model. By assumption, the output of the model approximately captures the conditional 3 -gram, which by Equation (72) equals | A X B |{| B | . Note that if | A X B |{| B | is non-zero, it must be at least 1 { T . This means, if the transformer is able to compute the conditional 3 -gram to within an additive error of 1 { 3 T , then Bob can simply threshold the output of the transformer to decide whether A X B ' H or not, thereby solving DIS p a , b q .

Since this communication protocol is deterministic, by Yao's lower bound in Theorem 7, the number of bits communicated between Alice and Bob must be at least m ' T { 3 ´ 1 . The total number of bits of communication in the protocol is 2 pH ` dp ` 1 (the last 1 comes from Bob having to communicate the answer to Alice), completing the proof.

E Lower bounds on representing k th -order induction heads: Proof of Theorem 6

In this section we prove the size-lower bound on attention-only transformers representing k th -order induction heads in Theorem 6. To enable this result to be better interpreted, we will break it down into two corollaries.

Corollary 1. Consider an L -layer attention-only transformer with 1 head per layer and relative position encodings, which satisfies Assumption 1. If L ď 1 ` log 2 p k ´ 2 q , the attention pattern in layer L of the transformer cannot represent a k th -order induction head.

Corollary 2. Consider an 2 -layer attention-only transformer with H heads in the first layer and relative position encodings, and assume that Assumption 1 is satisfied. If H ď k ´ 3 , the attention pattern in the 2 nd layer cannot represent a k th -order induction head.

We will first prove the result for the case L ' 2 and H ' 1 , which falls in the intersection of both of these corollaries. We will show that these models cannot represent k th -order induction heads for k ą 3 , under Assumption 1. We subsequently extend it to the general L -layer transformer (i.e., Corollary 1) in Appendix E.2 and to the general case with H ℓ heads in layer ℓ P r L s in Appendix E.3.

E.1 Lower bounds on 2 -layer 1 -head attention-only transformers

In this section we show that under Assumption 1, a 2 -layer 1 -head attention-only transformer cannot represent k th -order induction heads for any k ě 4 . We will prove lower bounds on the transformer when the input is binary, i.e., S ' t 0 , 1 u . With relative position embeddings, observe that the first layer of the transformer model learns representations of the form,

where note that the attention pattern only depends on n and i and not on x i or x n . These representations are input into the second layer, which realizes the attention pattern att p 2 q n,i , which is proportional

to,

We need this function to be maximized uniquely when x i ´ 1 ' x n , ¨ ¨ ¨ , x i ´ k ' x n ´ k ` 1 . Denoting ϕ p 0 q ' W p 1 q V Emb p 0 q and ϕ p 1 q ' W p 1 q V Emb p 1 q ,

where x 1 i Ð 2 x i ´ 1 . We can write this down as,

where M p 1 q n is a matrix of rank at most 2 and of the form,

which is independent of x 1 1 , ¨ ¨ ¨ , x 1 n . Likewise m p 1 q n collects all the vectors in the sum that don't depend on x 1 1 , ¨ ¨ ¨ , x 1 n . Now, observe that in the next layer, we wish to show that an induction head cannot be realized by att p 2 q n,i for each i ď n . We will show this for any value of i ď n ´ k .

In the second layer, we may write down the key vectors as,

Again, defining the vector m p 1 q i and the matrix M p 1 q i appropriately (having rank at most 2 ), this equals,

where y fi ' x 1 i ´ 1 ¨ ¨ ¨ x 1 i ´ k ‰ T and the vector m p 1 q i depends on x 1 i as well as the inputs x 1 i ´ k ´ 1 , ¨ ¨ ¨ , x 1 1 , which in this context, are treated as nuisance variables since they do not intersect with t x 1 i ´ 1 , ¨ ¨ ¨ , x 1 i ´ k u Y t x n , ¨ ¨ ¨ , x n ´ k ` 1 u . Henceforth we will avoid explicitly stating the dependency of m p 1 q i on the x j 's. Similarly, the query vector can be written down as,

where x m p 1 q n , | M p 1 q n and x M p 1 q n are defined appropriately, with x M p 1 q n and x M p 1 q n of rank at most 2 , and x is defined as ' x 1 n ¨ ¨ ¨ x 1 n ´ k ` 1 ‰ T . For an appropriate matrix M ˆ n,i , vectors m ˆ n,i and Ă m ˆ n,i and scalar m ˆ n,i , the dot-product of the key and query vectors can be written as,

Which is a linear function in x and quadratic in y , both of which lie on t˘ 1 u k . Note that the matrix M ˆ n,i has rank at most 2 since it is a product of M p 1 q i and | M p 1 q n , each with rank at most 2 . Next we introduce a lemma showing that if M ˆ n,i is inherently low rank, the quadratic form in Equation (90) which captures the dot-product between the key and value vectors cannot satisfy the property that for every y , the function is uniquely maximized at x ' y . In particular, this means that for any i ď n ´ k , there is some choice of x n , x n ´ 1 , ¨ ¨ ¨ , x n ´ k ` 1 such that there are x i ´ 1 , ¨ ¨ ¨ , x i ´ k such that for at least one j P r k s , x i ´ j and x n ´ j ´ 1 are not equal, but the attention score is larger than the case when x i ´ j were equal to x n ´ j ´ 1 for each j P r k s .

Lemma 2. If M ˆ n,i has rank ď k ´ 2 , it is impossible for f n,i p x , y q to satisfy the property that for every y P t˘ 1 u k , the maximizer is uniquely x ' y .

The proof is almost complete: if k ě 4 , then the rank of M ˆ n,i , which is at most 2 , does not exceed k ´ 2 . This means that when k ě 4 , any attention pattern realized in the second layer must satisfy the property that there exists a string such that the attention is no longer uniquely maximized when x n ' x i ´ 1 , ¨ ¨ ¨ , x n ´ k ` 1 ' x i ´ k .

Proof. For the purpose of brevity, define H k ' t˘ 1 u k . First consider the reparameterization,

Then, the dot-product of the key and query matrices can be written as,

Note that this function is linear in r x and therefore must be maximized on a vertex of the convex hull of the domain, Ă M ˆ n,i H k fi ␣ Ă M ˆ n,i h : h P H k ( . If M ˆ n,i has rank at most k ´ 2 , the rank of Ă M ˆ n,i is at most k ´ 1 and cannot be full rank. We show that this must imply that there is a vertex v P H k such that Ă M ˆ n,i v is not a unique vertex of the convex hull of Ă M ˆ n,i H k . This means that v cannot be a unique maximizer for r x when maximizing over all strings in Equation (92), and specifically y ' v is a witness to Lemma 2.

Below we discuss how to find such a vector v . Note that Ă M ˆ n,i is not full rank, which implies that there exists a vector n such that Ă M ˆ n,i n ' 0 . Without loss of generality, let n 1 be the smallest non-zero coordinate of n in absolute value. Then the vector n ´ 1 1 n has no non-zero coordinates in the interval p´ 1 , 1 q . We will show that sign p n ´ 1 1 n q is a good choice for v .

Consider two cases,

Case I. Every non-zero coordinate of n ´ 1 1 n is in t˘ 1 u . Consider any x P H k which matches with n on the non-zero coordinates. Consider x 1 which is the same as x , except a negation is taken on the coordinates where n is non-zero. Note that Ă M ˆ n,i x ' Ă M ˆ n,i x 1 , for the same value of x . This means that for any y . In particular, from Equation (92), both x and x 1 are maximizers, showing that Lemma 2 is true in this case. We circumvent having to find such a vector v in this case.

Case II. n ´ 1 1 n has non-zero coordinates which are not all in t˘ 1 u . In particular, at least one coordinate where this vector is strictly less than ´ 1 or strictly greater than ` 1 . In this case, observe that the sign vector r n ' sign p n ´ 1 1 n q P H k lies within, but is not a vertex of the convex hull of the set H k Yt n ´ 1 1 n u . The reason for this is simple to see when we assume that n ´ 1 1 n has only one coordinate which is not in r´ 1 , 1 s , say, the coordinate j ' 2 : here, r n can be written down as a convex combination (with non-zero coefficients) of n ´ 1 1 n and r n p 2 q ; the latter vector is obtained by flipping coordinate 2 of r n . When there is more than one coordinate not in r´ 1 , 1 s , we can peel away these large coordinates in n ´ 1 1 n by taking a convex combination of this vector with the vectors r n p j q for the appropriate values of j , to return the sign vector r n . Here, r n p j q is the version of r n where the j th -coordinate is flipped. This results in the following claim.

Claim 1. The sign vector r n lies within the convex hull of the points H k Yt n ´ 1 1 n u , but is not a vertex of this set.

In particular, we may write,

where α 0 ą 0 and ř n j ' 0 α i ' 1 . By left-multiplying this on both sides by Ă M ˆ n,i and noting that n lies in the null-space of this matrix, we get,

where note that ř j Pr n s α j is strictly less than 1 , since α 0 ą 0 . We may write this vector as,

Since α 0 ą 0 , this equation implies that the image of r n under Ă M ˆ n,i itself falls within conv p Ă M ˆ n,i H k q , but is itself not a vertex of this set. This means that r n can never be a maximizer of f n,i ` ¨ , y ˘ for any y , and in particular when y ' r n , thereby proving Lemma 2.

E.2 L -layer attention-only transformers with 1 head per layer: Proof of Corollary 1

Proof. The proof largely tracks the 2 -layer case, with the main exception that we keep track of how the maximum possible rank of the matrix M ˆ n,i grows as a function of the depth of the transformer. In the case the 2 -layer transformer, we show that it cannot exceed 2 . With the addition of more layers, we show that it cannot exceed 2 L ´ 1 .

Recall from the notation in Equation (85) that the output of the first attention layer is,

where M p 1 q n P R d ˆ n has rank at most 2 . Let us rewrite this as,

where M p 1 q n P R d ˆ T is causally masked to be 0 's when it operates on x i for all indices i ą n . Note that even with this causal masking, M p 1 q n has rank at most 2 , as discussed in Equation (85). By induction, assume that the output of the p ℓ ´ 1 q th attention layer is of the form,

where x 1: T fi ' x 1 T x 1 T ´ 1 ¨ ¨ ¨ x 1 1 ‰ T . Passing x p ℓ q n through the ℓ th attention layer, we get,

Define,

Then, we can write down,

We also inductively assume that for every i ď n ,

p i q M p ℓ ´ 1 q i has rank R ď 2 ℓ ´ 1 , and,

p ii q M p ℓ ´ 1 q i can be factorized in the form ř R r ' 1 u r ¨ v T i,r , where only the v i,r 's depend on i , but the u r 's do not depend on i .

Both of these conditions are true when ℓ ´ 1 ' 1 as evidenced by the structure of M p 1 q i in Equation (86) and noting that M p 1 q i is obtained from M p 1 q i by right multiplying by a diagonal mask matrix. Using the recursion in Equation (101), we prove that the induction hypotheses p i q and p ii q are true at layer ℓ as well. In particular using the decomposition in p ii q , observe that,

where for r 1 P r R s , u R ` r 1 fi W p ℓ q V u r and v n,r 1 fi ř i ď n att p ℓ q n,i v i,r . Since M p ℓ q n is the sum of 2 R rank 1 matrices and therefore has rank at most 2 R ď 2 ℓ , proving both parts of the induction hypothesis.

By induction, at the end of the p L ´ 1 q th layer, we have an output which looks like,

where M p L ´ 1 q n has rank at most 2 L ´ 1 . More importantly, note that by the causal masking, even though it appears to depend on the whole input sequence through x 1: T , note that x p L q n only depends on x 1 , ¨ ¨ ¨ , x n and not on the future inputs to this time n . In particular, by a similar argument as in the 2 -layer case (cf. Equation (85) to Equation (90)), for any i ď n ´ k we can decompose the dot-product of the key and query vectors at the L th layer as a bilinear form which looks like,

where x and y are defined as ' x 1 n ¨ ¨ ¨ x 1 n ´ k ` 1 ‰ T and ' x 1 i ´ 1 ¨ ¨ ¨ x 1 i ´ k ‰ T respectively, and M ˆ n,i has rank at most that of M p L ´ 1 q n , which is 2 L ´ 1 . In particular, if 2 L ´ 1 ď k ´ 2 , by Lemma 2 the proof concludes.

E.3 The general case: Transformers with H ℓ heads in layer ℓ : Proof of Theorem 6

The h th head of the first layer of the attention-only transformer learns patterns of the form,

where the last equation assumes a binary input sequence, defines x 1 i ' 2 x i ´ 1 and uses the notation ϕ h p 0 q ' W p 1 ,h q V Emb p 0 q and ϕ h p 1 q ' W p 1 ,h q V Emb p 1 q . We can further rewrite this as,

where each M p 1 ,h q n P R d ˆ T is rank 1 and applies a causal mask on the inputs x i for i ą n . Recall that the output of the first attention layer applies a projection matrix on the concatentation of r x p 1 ,h q n across h P r H 1 s and then adds a residual connection. The output can be written down as,

where,

p 1 q

Notice that the rank of the matrix M n is at most H 1 ` 1 . This is because the concatenation operation can increase the rank at most additively, and since each of the M p 1 ,h q n matrices are rank at most 1 .

Following through the proof in Appendix E.2 for the L -layer case, we can prove inductively that at any layer ℓ , the output looks like,

where the rank of M p ℓ q n is ś ℓ i ' 1 p H i ` 1 q . Invoking Lemma 2, if ś L ´ 1 i ' 1 p H i ` 1 q ď k ´ 2 , the attention-only transformer cannot realize a k th -order induction head at layer L .

F Model architecture and hyper-parameters

The experiments were run on one 8 ˆ A 100 GPU node.

G Additional experimental results

Assumption 1 suggests that the attention patterns att p ℓ q n,i in layers ℓ ' 1 , 2 , ¨ ¨ ¨ , L ´ 1 , as learnt by an L -layer attention-only transformers may only be a function of only the position indices n, i . In this section we run some additional experiments to test this conjecture. We train a 2 layer attention-only transformer with k heads in the first layer, on data drawn from a randomly sampled k th -order Markov process, and focus on the learnt attention patterns as a function of in the input sequence. Figure 7 plots the results of this experiment for k ' 2 and Figure 8 for k ' 3 . While in both cases there is some variance in the attention patterns learnt by the transformer in some of the heads, we believe that this is a consequence of the iteration budget of the transformer, and specifically the fact that even if the test loss appears to have converged, the transformer may still continue changing in the parameter space. Furthermore, when the attention patterns have some non-zero but small variance as a function of the input, a relaxation of Assumption 1, we also believe that the results we proved in Corollaries 1 and 2 and Theorem 6 should carry over approximately and leave this as an interesting question for future work. Conditional lower bounds of this nature, reliant on structural assumptions the transformer appears to demonstrate in practice are an interesting area of future research.

0

2

4

6

8

10

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

Sequence index

Attention weight

(a) First head

0

2

4

6

8

10

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Sequence index

(b) Second head

0

2

4

6

8

10

0.00

0.05

0.10

0.15

0.20

0.25

Sequence index

Attention weight

(a) First head

0

2

4

6

8

10

0.00

0.02

0.04

0.06

0.08

0.10

0.12

0.14

Sequence index

(b) Second head

0

2

4

6

8

10

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Sequence index

(c) Third head

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 paper presents experimental evidence about k th -order Markov processes being effectively learnable by small-depth transformers. We also prove theorems about the representation power of low-depth transformers, which are stated in the main paper.

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 state that our lower bound in Corollary 1 is a conditional result in Section 6.1, and that it is an important open problem to see whether the conditional statements can be removed. While our work talks about the representational power of transformers, we also mention that learning dynamics of gradient descent is an important direction that future work needs to address.

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: All theorems in the paper are stated and cross-referenced with proofs in the appendix.

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: Experiment hyperparameters are presented in Tables 2 and 3 and the code has been provided along with the supplementary material.

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: [Yes]

Justification: All the code has been open sourced; hyperparameter choices have been provided in Tables 2 and 3.

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: Tables 2 and 3 cover this information.

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: [Yes]

Justification: Standard error bars are provided on all plots.

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: Details in Appendix F.

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: No NeurIPS code of ethics were violated.

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: This work is a primarily theoretical study on the behavior of tokenization on toy problems (learning Markov chains). The societal impact of this research is not likely to be significant.

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: No models with a high risk for misuse were trained or released.

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: [Yes]

Justification: Code has been properly credited, via citing the relevant papers.

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: No new assets released.

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: No crowdsourcing or research with 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: No crowdsourcing or research with 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 attempts to provide a possible explanation the capability of transformer architecture for efficient next token prediction of (stationary) $k$th order Markov data. The main result is a constructive proof that a $3$-layer transformer with a single head per layer can emulate conditional $k$ grams, and it necessarily uses the non-linearities in the transformer architecture like layer-normalization. The paper also provides lower bounds on representational powers of the transformers for Markov data.\n\nSTRENGTHS:\n(1) The results presented in the paper serve as a crucial step towards understanding the prediction capabilities of transformers for $k$th order Markov processes with finite state spaces (of small to moderate sizes).\n\n(2) The paper provides a novel explanation of the (possible) way in which a layer-normalization (LN) might be used in modelling $k$-grams. To my knowledge, the non-linearities are discarded in typical prior work(s).\n\n(3) The lower bound is provided for single layer transformer (Theorem 5) is a strong evidence signifying of importance of multiple layers.\n\n(4) The paper is well-written and easy to follow.\n\nWEAKNESSES:\n(1) In spite of solving a general and interesting problem of generative modelling using Markov structure, the evaluation is very limited. For example, I believe that the paper does not have empirical validation of $k$-gram modelling by the transformer when the size of the state space $S$ is large.\n\n(2) The paper does not address or comment how the training data can affect the performance during testing. This step of training might be very significant when the size $S$ of the state space is high. It would be nice to see how the mechanism works when $S$ is of order of tens or a couple hundreds (my understanding is that Figure 2 is only for binary data, $S=2$).\n\nTypo in Architecture 1: $\\textbf{x}_i^{(1)} = \\texttt{Emb}(x_n)$ instead of $\\textbf{x}_n^{(1)} = \\texttt{Emb}(x_n)$\n\nQUESTIONS:\n(1) I believe that the state space is binary in your evaluations $S=2$. Can you please confirm this?\n\n(2) In Table 3 of Appendix F, the embedding dimension is mentioned to be chosen as per grid search over \\{16, 32, 64\\}. Should this be not the same as $d=6S+3$ as in the proof of Theorem 4 in Appendix C.3? In this case, if $S=2$, then $d=15$ is enforced. Can you provide the explanation for grid search?\n\nLIMITATIONS:\n(1) Although the result sheds light on the abilities of a transformer architecture in modelling conditional $k$-grams, this is only optimal for time-homogeneous stationary Markov data. Hence, the paper does not directly have a huge impact towards understanding how transformers perform on general stochastic sequences (like language tokens, or frames in a video).\n\n(2) The constructions in the paper work only for small state spaces (tens to a few hundreds), as the embedding dimension used in the proofs scales with its size.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the learning process and representational power of transformers on the data sequence generated by k-th order Markov processes (or k-gram). Theoretically, this paper proved that (1) an attention-only transformer with $O(\\log k)$ layers with one head for each layer can represent the conditional (in-context) $k$-gram. (2) Enhanced with MLP and layer normalization a 3-layer transformer can express the conditional $k$-gram task. They also complement their results with a lower bound for attention-only transformers/a 1-layer transformer with MLP and layer-norm with some assumptions. Detailed empirical validations are conducted to corroborate the construction (1).\n\nSTRENGTHS:\n1. This paper improved the construction using $k$-head, 2-layer in previous papers on the in-context $k$-gram data by constructing a $\\log k$-depth, 1-head transformer. The proof technique is based on the binary-tree-like aggregation procedure of the previous $k$ tokens information, which greatly improves the memory cost of the transformer parameters. Empirically, the construction can be partially validated by experiments showing that a 3-layer transformer can learn an in-context 8-gram. Also, a lower bound result is included (with some empirically validated assumption) to make this bound tight.\n\n2. The construction of a 3-layer transformer with the MLP and LayerNorm is quite novel. The technique is based on the unique ternary representation of integers. The ternary representation is to maintain the uniqueness of $v_i$ after normalizing the embedding vector. The role of the non-linearity in the transformer in this $k$-gram mechanism can be crucial for improving the representation power of the model, and this result can serve as an implication.\n\nWEAKNESSES:\nThe two main construction results improve the previous expressivity results. However, the second theorem (Theorem 4) still lacks empirical evidence showing that the constructed solution is the minimizer that Adam/GD converges to. The previous experiments only show that $\\log k$-depth transformer can somehow be learned. It is important to figure out whether the solution is only for construction or it can be actually learned in some way.\n\nQUESTIONS:\n1. Is it possible to add a set of experiments for the 3-layer transformer with LayerNorm and MLP trained on the k-gram model when $k>8$? Or showing that it is hard to obtain from optimization?\n\nLIMITATIONS:\nSee weakness.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the ability of transformers to learn $k^{th}$ order Markov chains. They first conduct experiments showing that transformers with 2 layers and 1 head can learn Markov chains of up to order $k=4$. Similarly, with 3 layers, they can learn Markov chains of order $k=8$. Based on these observations, they show theoretical results about the representation power. Specifically, they present novel constructions to show that attention-only transformers with 1 head can learn $k^{th}$ order Markov chains with $\\log_2(k+1)$ layers. On the other hand, 2-layer attention-only transformers require $k$ heads. These results show that increasing the depth is much more beneficial than increasing the number of heads. Next, the authors show that for the full transformer model, a constant depth of 3 suffices to learn $k^{th}$ order Markov chains with an embedding dimension of the order of the vocabulary size. This result reveals the benefit of non-linearities arising from layer normalization. They also present a lower bound on the depth requirement for 1-head attention-only transformers under some assumptions on the attention pattern learned by the model.\n\nSTRENGTHS:\n- Overall, the paper is well-written and easy to follow. The authors present nice intuitions about the constructions and the theoretical results.\n\n- The results in the paper showcase the benefit of depth over the number of attention heads in transformers, which contributes to our understanding of transformers. The paper also provides insights into the benefit of non-linearities that arise from layer normalization in transformers, which is interesting, and helps compare attention-only transformers with standard transformers.\n\nWEAKNESSES:\n- The main weakness is that some of the statements made by the authors about their results/observations contradicting prior work *lack clarity* and are not supported by a thorough *discussion of the potential reasons for the differences*. I think these statements should be accompanied by more context. Please see the Questions section for queries/suggestions regarding this.\n\n- There are some minor typos and grammatical errors that should be corrected. Please see the next section for details.\n\nQUESTIONS:\n- The statement in lines 28-31 needs further clarification:\n\n    - It’s not clear what range of $k$ is considered in prior works. It seems that the observations are not exactly in contradiction to prior work, since a) 2-layer 1-head transformers can’t seem to be able to learn Markov chains of order $k>4$, and b) the result in Table 1 and Theorem 2 states that for order $k$, the number of heads for 2 layers is $k$, at least for attention-only transformers. Table 1 is missing a column for the number of heads needed for standard transformers with 2 layers.\n\n    - Has prior work considered training transformers with 3 layers? How does the statement after (ii) contradict prior observations?\n\n- Regarding the experiments, particularly Fig. 2:\n\n    - In Fig. 2(a), the test loss gap seems a bit high. Can the authors include results for $k=1$ in this plot for a comparison?\n\n    - Can the authors share results when training the 3-layer 1-head transformer on Markov chains of higher order $k>8$? Is there a reason why higher-order Markov chains were not considered for the experiment in Fig. 2(b)? I am curious if there is a gap between the theoretical result about the representation power (which holds for $k>8$), and the training dynamics.\n\n- Other suggestions:\n\n    - In the abstract, it would be good to emphasize the word ‘constant’ in line 14.\n\n    - I suggest using $V$ instead of $S$ for the vocabulary size.\n\n    - It would be good to include some discussion on whether similar construction techniques (for Theorem 3) have been used in the literature.\n\n    - I suggest including some discussion on related work on the role of softmax attention, the benefit of layer normalization techniques, and the benefit of depth in neural networks.\n\n- Typos/grammatical errors:\n\n    - Missing citation in line 23.\n\n    - Extra ‘studies’ in line 27. Extra ‘them’ in line 132. Extra ‘how’ in line 272.\n\n    - Please check the phrasing in lines 220-222.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the representation capability of transformer with different number of layers or heads when learning k-th order Markov Process. They provide theoretical results demonstrating that attention-only transformers with O(log2(k)) layers can represent the in-context conditional k-th order Markov Process. This conclusion is supported by empirical results and is novel compared with previous training dynamics results. Adding LayerNorm, they prove that standard transformers with just three single-head layers can represent arbitrary order Markov processes. The paper also presents lower bounds on the size of transformers needed to represent certain attention patterns.\n\nSTRENGTHS:\n1. The writting is good and easy to follow. And the proof is relatively solid. Most of the intuition is clear, like using hierarchical way to construct the intermediate logits which include information of multiple tokens.\n2. The main contribution of this paper is that the authors prove there exists some transformers which have fewer head and logarithm-limited layers that can learn the determined k-th Markov Process. And these results are supported by the empirical experiments, and it's more related to the real-world case. (Like we don't need too many heads)\n\nWEAKNESSES:\n1. This work analyzes the representation capability of transformer of learning Markov Process, however, there is no clear statement of error bounds in the key theorem like Theorem 4 (although there are some helpful discussions like Remark 2). And this may related to some questions like Questions-2.\n2. The intuition of the proof of Theorem 5 is still not very clear, maybe adding more examples will be helpful, like how LayerNorm works.\n3. Like mentioned in limitation parts, it's more focused on representation capability rather than training dynamics. The latter one may be more difficult for analyzing. More questions in Questions-1.\n4. In Figure 2, the authors should better add k=8 and k=16 to further support the conclusion.\n5. There is a citation typo in Line 23.\n6. Maybe you can explicitly say that learn the conditional k-gram model is equivalent to letting transformer generate the same logits as that model (i.e., Eq.1), otherwise the reviewer maybe confused by the conclusion in the main theorems.\n\nQUESTIONS:\n1. This work is more related to representation capability of transformer, rather than training dynamics, which are the focus of related works [6], [7] mentioned in the papers (Line 114 - 133), so is it possible that the construction of k heads is more theoretically friendly than the that in your paper for analyzing dynamics like gradient flow? And what's the possible way for training dynamics under your structure?\n2. It seems that you still need $\\Omega(k)$ bit precision for learning $k$-th order Markov Process, but in real-world case we don't need such a high precision. Is this necessary? and is it required by the constant layers transformer setting and not required by the $O(log(k))$ attention-only setting?\n\nLIMITATIONS:\nThe authors have mentioned their limitations in the end of paper, like they just focus on representation capability rather than training dynamics, which maybe the futural direction.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the representation capacity of transformers in in-context learning of order-$k$ Markov chains. First, the authors theoretically show that $O(\\log (k))$ layers are sufficient to represent $k$-th order induction heads in attention-only transformers. The paper also demonstrates the benefit of non-linearities, such as layer-norm, by showing that a slightly modified transformer (modified from the original architecture) with constant depth is also sufficient to represent the in-context conditional distribution.\n\nSTRENGTHS:\nGiven the prevalence of in-context learning, I believe understanding the ICL of Markov chains is a very important problem, not only because of the Markovian nature of language but also the simplicity and control it provides in a synthetic setup. This is a good work that builds on recent research on ICL of Markov chains by showing that in a slightly modified transformer, scaling of heads with the order of the chain ($k$) is not necessary to represent $k$-th order induction heads. The lower bound (albeit contingent on the $k$-th order induction head assumption) also provides useful insights. I really enjoyed reading the paper; it was very clearly written and easy to follow. Generally speaking, all the ideas are laid out very clearly, with intuitive explanations that follow before and after the theorems.\n\nWEAKNESSES:\n1. It would be beneficial to compare single head results with multiple heads (e.g., in Figure 2a, $2$ heads for $k=2$ and $4$ heads for $k=4$) as a sanity check, given previous works [1] that highlight the multi-head requirement for learning order-$k$ induction heads. I mention this because, to me, it seems like in Figure 2a, the loss isn’t approaching zero unlike in Figure 2b.\n\n\n2. The work mentions \"long training\" to observe contradictory results (2-layer single head tfs being able to learn up to order-4 MCs in-context). How long are you training compared to [1]? How significant is this effect, and what impact do you think it has in terms of the underlying mechanism being picked up? I think this needs to be discussed further.\n\n\n[1] Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The\nevolution of statistical induction heads: In-context learning markov chains, 2024.\n\nQUESTIONS:\n1. What do you think is the reason 2-layer, 1-head transformers can learn up to order-$4$ MC (if they do, see Weakness 1)? What mechanism do you think is inherently being picked up? Is it learning order-$k$ induction heads or some alternate mechanism? I ask because if there is an alternate mechanism, it could be useful for the lower bound (Thm. 6) in establishing if \"the lower bound representing $k$-th order induction heads implies an unconditional lower bound.\"\n\n2. Regarding Theorem 4, even though the embedding dimension requirement is reduced compared to attention-only transformers (Thm. 2, Thm. 3), it still feels quite large compared to what we see in practice. Any comments on this? How tight do you think this is?\n\nLIMITATIONS:\nYes, the authors have discussed the limitations separately in the Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-19"}