Transformers on Markov data: Constant depth suffices
Authors: Nived Rajaraman
Marco Bondaschi
Ashok Vardhan Makkuva
Kannan Ramchandran
Michael Gastpar
Venue: NeurIPS 2024 poster (2024)
Keywords: Depth
k-gram
Transformers
Representation
Primary area: interpretability_and_explainability
TL;DR: Low depth transformers with $1$ head per layer can learn $k^{\text{th}}$-order Markov processes

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
Nived Rajaraman
Marco Bondaschi EPFL
Kannan Ramchandran UC Berkeley
Michael Gastpar EPFL
UC Berkeley
Ashok Vardhan Makkuva
EPFL
Abstract
Attention-base based transformers have been remarkably successful modeling gen- erative processes across various domains and modalities. In this paper, we study the behavior of transformers data drawn from -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 empir- ically which contradicts previous findings: when trained for sufficiently long, a transformer with a fixed depth and head per layer is able to achieve low loss sequences drawn from hth -order Markov sources, even as & grows. Fur- thermore, this low loss is achieved by the transformer's ability represent and learn the in-context conditional empirical distribution. On the theoretical side, our main result that transformer with single head and three layers can represent the in in-context conditional empirical distribution for -oddr Markov sources, concurring with our empirical observations. Along the way, we prove that attention-only transformers with ((log((k)) layers can represent the in-context conditional empirical distribution by composing induction heads to track the pre- vious & 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 Markov sources. Code is available at: thub. com/Bond1995/ fonntantt.. consstattdepthht Transf ormers.
https: //githu.
Introduction
irention-based transformers have revolutionized the field of natural language processing (NLP) [1 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 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 -order Markov process, where the next symbol's conditional distribution depends only on the previous symbols. Recent studies [7, 6, 8] have investigated the ability of transformers 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 bit concerning real world data generating processes often contain long-range dependencies. How is that transformers succeed 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 thh-order Markov sources?
*Corresonddence to nived. wwwreddrajaraaaaee rrajaraman@berkeeeyy.. edu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Attention-only
Standard
L
[loga (k + 1)1
3
2 0
0
Xn+1
H
1
Figure 1: hthh -order Markov process for k 4. The symbol X is sampled from the distri- bution P(-[Xn, xn-1, 2n- Xx-3) which only depends on the last 4 symbols (marked in red).
Table I: Each column in this table indicates that there transformer with T. layers and H heads in the first layer which can represent the conditional k-gram model.²
To understand the nature of this phenomenon, we train low-depth transformers k Markov sources. These experiments result in two surprising empirical phenomena that seem to contradict previous findings: when trained for sufficiently long, (i) 2-layer, 1-head transformer kth-order Markov processes for le as large 4, (ti) a 3-layer, 1-head transformer is able to achieve low test loss on sequences drawn from 1tth -oddr Markov sources, even grows to be as large as (Fig. 3). In both cases, the values of k for which the models appear to learn hthh -oderr 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 odh-order Markov processes is not complete and raises broader question:
learn
What the interplay between depth, number of heads and non-linearity in learning tt 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 follows:
1. We show, rather surprisingly, that the standard transformer architecture with layers and 1 head per layer is capable of representing the conditional k-gram model (Definition 1), and thereby learn Athh -order Markov models in-context.
2. Along the way to building up to this result, we consider the simpler family of attention- only transformers and show that they can represent the conditional k-gram model with [log2(k + 1)] layers.
3. Under natural assumption on the nature of the attention patterns learnt by the transformer, then argue that for k >/2 3 attention-only transformers need at least [1 loga (A 2) layers to represent a "1th induction head" (Definition 2). Empirically, transformers are observed to learn 4t ooder induction heads whenever they achieve small test error [6].
The last result is a consequence of more general tradeoff between the number of layers, L, and heads layer, H. an attention-only transformer requires to represent ht -oddrr induction head, under natural assumption on the learnt attention patterns. conjunction, these results also reveal the role of non 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 4h 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 Opxq (resp. 1px;) refers to the all-zero (resp. all- one) matrix. When
clear from the context, we omit the dimensions of a matrix. Define [$] + (1.2..... S} for S e N. T(-) denotes the indicator function and Unif(S) denotes the uniform distribution over set
1.1 Related work
There is a large body 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
"The requisite embedding dimension and bit- precision to achieve target additive approximation is discussed in more detail in Sections and

---

transformers, and in co-ttnt learning, [13, 14, 15] study the representation capabilities of transformers and show properties such as universal approximation and triingg-omleeeeesss... 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 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 different phases of learning.
Preliminaries
We provide the necessary background for Markov processes, the conditional k-gram model, and the transformer architecture.
2.1 Markoy 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 only influenced by the most recent states. More formally, sequence (Xn)n31 is 1th Markov process a finite state space [S] with the transition kernel P. if surely,
P(Xn-1 XI, X..) = P(Xn+1 [n- Xn)
This property allows us to capture the conditional distribution at any position using only its previous symbols. This motivates the notion conditional k- k-gram, its empirical counterpart, defined for any sequence "ii....,TT...
Definition (Conditional k-gram model). Given a sequence (z1, ,n) of length n in [S]", the conditional k-gram model Pru(- "1,"". , In) corresponds to the in-context estimate of the distribution over symbols conditioned on the last l symbols, i.e. for x E [S].
Prx I1, ,nn) 2i-k41 1(24 84-1
Fn ................................ Fi- kn-k+1)
21-k41 [ Tn.....
k+ 11
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 prior distribution [27].
In our experiments, we will consider 4thh -oder Markov kernels sampled from Dirichlet prior with parameter 1. Namely, the transition P(- ",, _x. ig) is sampled independently and uniformly on the S-dimensional simplex A,, for each tuple (f1, (k).
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 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 sequence of length T is defined in Architecture 1.
The attention scores in layer L, {a & n). are computed Softmax(( ((WWl)) (ax)) +
Pn-j
[n]]). The superscript (f) indicates the layer index, and the matrices (8) ( dddd capture the key, query and value matrices in layer É. Note that the attention-only transformer may include feedforward layer with linear activations, i.c. linear
().KK
(e) (2)
w
K.
3

---

Pr(
Linear + Prodiction
FFN
Input:
2
0
3
0
3
Layer
Layer norm
2:0
match final symbols
Layer 8
Uniform
distribution
Multi head POS
Layer
(a) Conditional g-gram model. The conditional k-gram is the in in-context estimate of the Markov process and realized two steps. The first step find the locations in the sequence (marked red) which match the final & symbols (functionally, tion head). The conditional k--rrmm model returns the uniform distribution the next symbol these loca- tions (marked blue).
Layer norm
2
Emb
Emb
Embh
-order induc-
Transformer layer
(b) Transformer architecture. relative position encodings.
POS refers the
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 blue), which translates the key and value vectors depending the relative position of the embedded symbol.
for 1,2, T; (1) Emb( E Rdd.
(Input embeddings)
for = 1, 2, ... ,,, do
for 1,2, ... ,T, do
[1(() [i][n] att (0) n.
wil) (z! ( CRd
(Attention)
1 ((+1) x + (())
(Residual)
Here, att Sof toftmax;
W (z: + P (2:K), w(
x > :ie [s])).
logity
An (L+1) (L+1)
R
(Linear) (Prediction)
Pro (- T1, (r) = f (logit,) e R.S.
Architecture 1: Attention-only transformer.
The extension to H heads is straightforward, where in each transformer layer there are H attention (E,H)
layers in parallel, resulting in (2.1) .... yn E Idd for each n. These vectors are concatenated and passed through a linear transformation W RdH - Rd 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$ and the resulting vector is passed through a non-linearity f, usually a softmax, result in the model's prediction of the next symbol probabilities. The theoretical results in this paper will choose f = ReLU(-).
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 learn order Markov pro- cesses. The experimental results present in the literature suggest that in order for a 2 layer transformer model be able learn randomly sampled Markov process of order 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 A:th -oddrr Markov processes in-context. Similarly, the
crucial the number

---

Test loss gap from the optimal
Test gap from the optimal
0.4
0.2
0.1
0.0
5000
10000
15000
0
5000
10000
15000 20000
25000
Iteration
Iteration
Figure 3: Gap with the optimal test loss for (a) 2-layer, 1-head transformer model (above), and (b) a 3-layer. 1-head transformer (below), averaged over 3 runs for each k. The models learn the conditional k-gram model for randomly sampled k-th order Markov processes, even for large
authors of [8] study generalization of learning oother Markov processes to learning causal processes on degree k graphs. The theory and experiments pertain to 2-layer head transformers.
In Figure 3, we train 2 and 3-layer transformers with single head data drawn from random Markov processes of various orders drawn from a Dirichlet prior. With 2 layers and a single head, that the model is able to learn even order-4 Markov processes, and go beyond the simple order-1 processes which were projected to be the limit of ability to learn. Likewise, with 3 layers, transformers are able go much further and learn order-8 Markov processes, which was the largest value we evaluated on. These results contrast with our current understanding of how induction heads realized in the parameter space [6, 8] existing constructions which realize these attention patterns require & heads when the number layers is 2, and it's unclear how to implement them with fewer heads. At high level, each of the & heads play a critical role where, loosely speaking, the hled looks back 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.
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 orddrrder Markov processes whose parameters are drawn from Dirichlet prior. The model is observed to learn 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 (1) T(i
pattern att n, at position n 1 in the embedding vector at time n. the second layer, the attention layer picks out those indices where _n-1 IT, the final symbol in the sequence. At these positions, since In-1 27, one would expect that the next symbol vn is a good predictor of #T and the model uses this information to predict the next symbol #T+1 according to its conditional empirical estimate, Pri (#T 1,,, ... ,IT), i.c. the conditional 1-gram model.
1). thereby allowing the model to capture information about the symbol
Theorem The conditional 1-gram model can represented by 2-layer and 1 1-head attention-only transformer with embedding dimension d 3S 2.
Although a version of this result is proved in [6]. we include a proof in Appendix A for completeness.
Remark Theorem and other results to follow, we de-emphasice the role of the bit bit-precision to which the transformer is implemented. That said, note that when the constructions in Theorems I are implemented (loo(TT) bits of precision, the representation results are realized up to an additive O(1/T) error.
The ideas Theorem readily extend to representing the conditional k-gram model, by instead
using & heads in the first layer. The jth head learns the attention pattern att II(i n - j) and
5

---

concatenating the outputs of the heads, the model learns to aggregate information about 2n. kn in the embedding vector time n. The second layer realizes what is best described as a "kth "khe- induction head, where the model learns to pick out those positions n where for every
[k]].
Jn-j -r-j+11 i.e. the history of length sequence see Figure 4). This mechanism is also referred to as a long-prefin induction head [28].
those positions match the final k symbols in the input
Definition (Higher-order induction head), A 1 1-head attention layer is said to realize a -od-rr in- duction head if on any sequence (zi,-.. ; IT) [S] for any fixed n. & T, a function of the input sequence, alIn,T is maximized if and only if un-j = TT- j+ for every j E [k].
otddr induction heads generalize the concept of an induction head [20], and keep track of the positions & 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 temper- ature within the softmax of this attention layer results in an attention pattern which converges to the uniform distribution over those positions where the final i symbols J - , TT are seen previously in the sequence. Loosely, this allows the model to "condition" on the last & symbols in the sequence. With k heads, the model can ag- gregate information from the previous is po- sitions and implement a hthh -oder induction head, which leads to the following result. A full proof is discussed in Appendix A.1.
Input: 1
0
3
0
1
0
1
2
4
5
6
7
10
attr.n
0.5
0
2
4
5
6
7
9 10
71
Figure 4: 1thh -ordr induction head for A: 2. The at- tention pattern attr, is maximized for those values of n at which 2r-j+1 In-j for all E [A]. These the positions where the k-length prefix at those posi- tions matches with the last k symbols in the sequence.
Theorem 2.
The conditional k--ramm model can be represented by attention-only transformer with 2 layers, k heads and embedding dimension = (k + 2)S k
While this result is positive, it suggests that 2-layer transformer requires approximately k times as many parameters to able to represent the conditional k-gram model. The first result we prove that increasing the depth of the model is exponentially more beneficial, in that a transformer with (log(A:) depth can estimate in-context k-grams.
Theorem 3. The conditional k-gram model can be represented by atlention-only transformer with relative position encodings, with = [log, + 1)] layers and head per luyer. The embedding S.
dimension is & 2k(S +
With layers and & heads, the transformer aggregates information about each of the previous positions one step at a time through the & heads. However, with ((log(k:) 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(2) is chosen linear combination of x!! Emb(x,)) and (1) Emb(#nn1). This allows the embedding position n to aggregate information about 4n and Tn-1- In the same vein, in the second attention
layer, the model aggregates information from (2) and (=(2) in x the former has information about Fn and r,-11 and the latter has information about Fn-2 and .n-3- This expands the "window" of x; 's on which wn depends on to size 4. In the 4th layer, the model aggregates information from x!! and x. (6) n222 which allows wn (f+1) to effectively depend on the x; 's window of size 26+1 starting at position n, namely En.... - In the final layer, the embedding at position i, x (1) for L [log2( (k 1)] depends on En. In-1, In-k- In the last layer, the model can realize the dot-product (W()) (1.) (1. w(( ET プソー I(Zn;
-+1) by choosing the key and query vectors appropriately. By increasing the temperature in the attention softmax, the attention pattern realized is the uniform distribution values of such that Σn-j Ø7--+1 for every E [k], i..., order induction head. The full proof of this result is provided in Appendix B.
(L)

---

atly, or exeep
(realizes a & do"-ordrr induction bead)
Layer
Layer
Emb(Zin)
Emb((r)
FFN
ay
Layers +
Emb(
Emb(rr))
Multi-head + POS
Emb
ar
21s
Rearranged transformer layer with layer normalization FFN.
(b) Realizing a -order induction head in 3-layer transformer Following the architecture in Figure Sa.
Figure 5: Disassembling the constant depth construction. The first two layers are critical in the model's 's ability to capture information from the previous k positions, Layer normalization plays critical role in the 3rd layer which realizes a -order induction head.
While this is promising step toward understanding the behavior transformers exhibit in Figure showing that depth plays an important role in their ability to represent conditional k- b-gram models, the picture is still not complete. The experimental results in Section that a transformer might not even require logarithmic depth to be able to learn --h-order Markov processes approximately.
not preclude possibility
the next section, we will study constant-depth transformers and establish 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 = n 1, n 2, ... which the model has access to information about at the current time n. By composing L 22 k:)) attention layers, the model is able to collect enough information within the output embedding xn (L+1) to be able to realize a -th-order induction head in the next layer. In this section, we prove that adding non- coneeiterr the architecture, in the form of layer normalization, can significantly change the mechanism in which the transformer realizes this 1 -order head. In particular, there are constant depth architectures which allow -dd--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, 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 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 attention- only design presented in Architecture is the addition of layer normalization and feedforward layer in the for-loop over e [T] for each transformer layer L. The differences between Architectures 2 and 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(S).
Remark 2. Although the proof stated does not bound the approximation error arising from finite bound on the bit precision of the transformer, theory, it should suffice have 2(log(T) + k) bits per parameter for the statement of Theorem 4 to go through with an O(1/T) additive approximation error. The main point that none of the weights of the model exceed exp(k) and with log(T) additional bits per parameter, the approximation error scales O(1/T).
7

---

// x Eveln att W (z) + p
RR
(Attention + Residual))
yn 1 ReLU (w E R
(FFN)
(1) plaxl E Rd
(0 yn
(LN)
2 (1+1) (00 + a(() R2
(Residual2)
Architecture Modified transformer architecture. The computations above carried out for each n E [2'] in each layer & E [4]. In the layer normalization step (LN), the feature mean is defined as, Unif((((( (e
and the feature variance o² En-- Kef (2) 91 x
5.1 Proof sketch
In the attention-only transformer with 2 layers and & heads, the model is able to keep track of where the final symbols in the sequence appeared previously (i.e., a -orde induction head) by, loosely, using each head keep track of the occurrences of one of the final & symbols. On the other hand, with the benefit of more depth, with L = (llog(k) layers, the model is able to collect enough x! (L+1) be able realize the same behavior: However,
information within the output embedding 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 -order induction head keeps track of the indices such that Vj E [A]] Ti-s =
Tn-j+1- Defining 2; + - 2'ex, 2'ea;-341 j41' notice that the condition {Vj [k:], Zi-j In-j+1] can equivalently be captured by writing (z,-1 z, 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, atnn,--1 2' for sigh 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, which is maximized whenever 2 Zn.. While dot-product attention naively encourages those values of for which 2;-1 and 2,, 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 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.
2;
Zn
att o exp
(1)
1
where is the temperature parameter. Then, as the temperature grows, the attention pattern essentially focuses on those values of for which z,/1lzi-1|z zn/ z With this attention
pattern, we are thus very close to the statement we wanted to check, (z,-1 z,). As it turns out, for the special structure in the 2,'s considered (dyadic sums of one-hot vectors), may write down,
2 Zn
2 1z = /n/[|zn.
A quantifiable equivalence is provided in Lemma 1.
Realizing Lg-norm attention (eq. (1)).
Observe the equivalence,
2i
Zn
2
Zn
2
Taking a softmax on both sides, notice that the RHS (up to an additive constant) is the Lg-norm based attention, while the LHS is the usual dot-product attention between 2/-1/||z;all and 2 z,1l2- Thus on normalizdd vectors, Lo-morm attention and dot product attention are but the same.
While the first layer of the transformer computes the z,'s by weighted summation, layer normaliza- tion fills the last missing piece of the puzzle which normalize them to unit norm. This is consequence of defining the embedding vectors appropriately, discuss more in Appendix C.I.

---

From this step, realizing the actual conditional k-gram model follows readily. In particular, as the temperature A. in the attention grows, the attention pattern zooms in on indices e In + {k & i & n Vj [i] = In-j+1] in the last layer. The value vectors at this step are the one-hot encoding of 24; putting everything together, the logits realized by the transformer are,
(oTit)(DT+1)
Σ iEZ,
I(x; =T),
[nnl
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 slightl and has an additional layer. The first two layers of the transformer respectively compute 2; and 2;-1 which are added to the embedding vector at time i. This is
important because we need to test whether 2 _ nn and not whether 2; _ Zn or Zi-1 Zn.
Summary. The construction can be summarized as follows: the first layer computes 2n - 20-1 Cxa-, by choosing appropriate value vectors and relative position embeddings to realize the attention pattern attm,,n-; oc 2'1(1 & i & k). The layemnrr that follows subsequently can be replaced by RMSnorm, by simple trick which we discuss in Appendix C.I, resulting in zn/ 5n/||zn] l2 to be appended to the embedding at time n. Using a very similar construction, layer 2 computes žn-1/ žn-1 1122 which is added to the embedding at time n. Finally, in the last layer, the dot-product (2-11/1=,-11/ /n/l (xnll2> defines the attention score, and as the temperature x grows, the pattern converges to Un.. 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 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- lear-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 able compute the conditional 3- 3-gram on inputs drawn from {0, 1, 2 to within an additive error of 1/3T, then 2pH + dp + 2 > T/3.
Choosing the bit precision to be = (log(T))) this implies that for transformers with 1 layer, the sum of the number of heads and the embedding dimension must be at least S2(T/ log/)) in order to represent conditional 3-grams to within an additive error of 1/3T.
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-oniy transformer with O(logz()) (k)) layers and 1 head per layer can represent conditional k-grams on its input sequences. With the addition of lineaiicas,sies, Theorem 4 shows that the model can represent conditional k-grams using just 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 attention-only transformer with layers and head, order-1 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 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 [H] realize attention pattern where att,, (t.h) only depends
the positions and and on & and h, but not on the input sequence "1,""" ZT.

---

seight
Index 3
Indox
15 Index
Indes
Sequence index (i)
Figure 6: Attention matrix of the first attention layer, for a 2-layer 1-head transformer model trained on an order-1 Markov process, averaged across 100 input sequences of length 128. (a) and (b) plot the mean and standard deviation of the first 32 rows and columns of the attention matrix, while (c) zooms in on the column for almost all indices the attention layer focuses only on the previous symbol 2n-1. (b) shows that the attention pattern does not vary much with the input sequence considered, thereby providing evidence toward Assumption 1, More discussion in Appendix G.
10 and plots the mean attention for this column. (a) and (c) show that
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 hthh -oderr induction head (Definition 2) in the last layer. Although learning a hthh -oderr induction head is not strictly necessary for the transformer to be able to represent conditional grams, note that every construction we have considered so far (cf. Theorems to 4) 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 extension of the -"horder induction head to general causal graphs. Our main lower bound is the following result.
Theorem 6. Consider an L-laaer transformer with be heads in layer L. Assuming the transformer L-1 satisfies Assumption 1, if t=1 (He 1) k -2, 2, the attention pattern in layer L cannot represent -order induction head.
While this lower bound is not unconditional, meaning that it does not directly imply that the trans- former 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 significantly different mechanism to realize the conditional k-gram model.
Theorem 6 implies that under Assumption 1, 2-layer attention-only transformer with head cannot realize a -ohhorder induction head for any >/ 4. Likewise, under the same assumption, a 3-layer attention-only transformer with head cannot realize 1eth g-h-order induction head for any k: > 6. These results give more weight to the experiment in Figure where we observe that a 2-layer transformer learns -order Markov process for & = 4 and 3-layer transformer learns - Markov process for k = 8, and show that non-lineasisies in the architecture allow the transformer to break past the size barriers in Theorem
7 Conclusion
We observe empirically that 2 and layer transformers are able to learn -odder Markov chains for much higher values of k than previously anticipated. We show there are ((log()))-layer constructions of iatention-only transformers which are able to learn the conditional l- k-gram model, which the in-context MLE of the Markov model. With non-linearities in the model, show that 3-layer 1-head transformer is capable of representing the same. We show that 1-layer transformers cannot represent conditional k-grams for any l unless the number of heads embedding dimension scale almost linearly in 7. We also prove a conditional lower bound on the depth and number of heads of attention-only transformers to represent order induction heads, under assumption the realized attention patterns.
Aeknowledgments and Disclosure of Funding
The work was partially supported by NSF Grant CCF-2211209 and Swiss NSF Grant 200364.
10

---

References
[1]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998-6008, 2017.
[2] 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-19011 2020.
[3] 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.
[4] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners, 2021.
[5] 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.
[6] Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024.
[7] 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.
[8] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent, 2024.
[9] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11080-11090, 2021.
[10] Angeliki Giannou, Shashank Rajput, YonY Sohn, Kangwook I Lee, Jason D. Lee, and Dimitris Papailiopooloss. Looped transformers programmable computers. In Proceedings of the 40th International Conference on Machine Learning, pages 11398-11444 1144, 23-29 Jul 2023.
[11] 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.
[12] 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.
[13] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of equeence-t--squeeee functions? In International Conference on Learning Representations, 2020.
[14] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention Turing-complete. Journal of Machine Learning Research, 22(75): 2021.
[15] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating Turing machines with transformers. In Advances in Neural Information Processing Systems, volume 35, pages 12071-12083, 2022.
[16] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transform- learn shortcuts to automata, 2023.
[17] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal, On the ability and limitations of trans- formers recognize formal languages, 2020.

---

[18] Clayton Sanford. Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth, 2024.
[19] Clayton Sanford, Daniel 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 Neural Information Processing Systems, volume 36, pages 36677-36700 Curran Associates, Inc.. 2023.
[20]
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.
[21] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Advances in neural information processing systems, 35:24824-24877 2022.
prompting elicits reasoning in large language models.
[22] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.
[23] Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning, 2024.
[24] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2023.
[25] Jesse Hoogland, George Wang. Matthew Carrugia-Roberty Liam Carroll, Susan Wei, and Daniel Murfet. The developmental landscape of in-context learning. arXiv preprint MxXiv:2402.02344 2024.
[26] Nived Rajaraman, Jiantao Jiao, and Kannan Ramchandran. Toward theory of tokenization in Ilms, 2024.
[27] J. R. Norris. Markov Chains. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1997.
[28) Nicholas Goldowsky- Chris MacLeed, Lucas Sato, and Aryaman Arora. Localizing model behavior with path patching. arXiv preprint arXiv:2304.0599 05969, 2023.
[29] Andrew Chi-Chih Yao. Some complexity questions related to distributive computing (prelimi- nary report). Proceedings the eleventh annual ACM symposium on Theory of computing. pages 209-213, 1979.
[30] Matteo
Pagliardini.
GPT-2
modular
codebase
implementation.
http://ggittbb.co//ppfm////l haspsp/ppmimmilmmmmblllleeee 2023.
12

---

Appendix
Table of Contents
A Proof of Theorem
13
A.I Extension to k-heads: Proof of Theorem 2
14
B Proof of Theorem 3
15
C Proof of Theorem
18
C.I Modifying the definition of layer normalization
18
C.2 Notation and supplementary lemmas
18
C.3 Proof of Theorem 4
19
D Representation lower bounds for 1-layer transformers: Proof of Theorem 5
22
E Lower bounds on representing -order induction heads: Proof of Theorem
23
E.1
Lower bounds on 2-layer 1-head attention-only transformers
23
E.2 L-layer attention-only transformers with head per layer: Proof of Corollary
26
E.3 The general case: Transformers with Hy heads in layer & Proof of Theorem 6
27
F Model architecture and prrre-paramesers
28
G Additional experimental results
28
The notation of E Rd' refers to the one-hot encoding of i in d' dimensions. In other words is the 7th standard basis vector in d' dimensions. The notation Blkdiag((AA) Ar. Amm)) refers to the block diagonal matrix with 7th block as A,
Notation.
A Proof of Theorem 1
We will first prove Theorem In the first layer, choose the embeddings as,
a (1) Emb(x,,) = [1x2 e 0 01x2s] RR.
for a constant ń > O to be chosen later and = 25 + 2. The relative position encodings will
essentially be supported on the first two coordinates, the middle S coordinates
one-hot encoding
of the symbol na and the last 25 coordinates are 0. The relative position encodings in the first layer
chosen to be Pn-i (1).K H (-1 + I(n i 1))) of IRd and ()).V 0 e IRd. Choose w (1)
and (i()) (1) to be (et)) Rdxxd With this choice,
(l:))) w ((1)))
kIl(n i 1)
(5)
As - s0, the attention pattern (which takes the softmax over of these inner products over i [n]) computes,
att (1)
ll(i
1)
(6)
for any n > 1. Choose the value matrix as,
((2+5)
0
0
wa
(1)
Isxs
0
E Rddd
(7)
O
0
And with this choice and the residual connection, get, ex
n
(2)
11x2
0 E d
(8)
which serves as the input to the 2nd transformer layer.
13

---

Vi(2) In layer 2, the relative position encodings Pn-i and Pa-i are all set O. The key matrix picks out the exx block out of æn In particular, these matrices are chosen so that,
K.(2)
Layer 2.
(2) and the query vector picks out the ex block out of x 2-1'
w (2)
K x [11x
S
oJ' E R
2111
1( (2) d = [11x
0 E
Taking the inner product of these vectors, and taking x - s0, observe that the attention pattern concentrates on the uniform distribution over all coordinates i such that 24-1 2n. More formally, the attention pattern for any n > is,
I( In) Li=2 1(x4-1 n..))
att
(10)
assuming Σ 11(2,-1 = zn) > O. Having realized this attention pattern, may choose the value and subsequent linear layer appropriately. The value matrix simply picks out the ex, block from (2) and places it into the last S coordinates of x; (after scaling down by factor of н), realizing the logits,
(3) and the linear layer simply extracts this block and outputs
logit,
Σ 1(a-11 xn). cr.
(11)
Ex= - = Ind
if Ei= 1(x1-1 2n) > O. In particular, under the same condition,
logit,( 2n=2 I(rn IT+1, Fn-1 = TT)
(12)
, 2T)
assuming Ee= I(zn-1 = =TT, which is the conditional 1 1-gram model.
A.1 Extension to k-heads: Proof of Theorem 2 In the first layer, the embeddings are chosen to be,
Emb(x,)) K. 01xk
01x(*+1JS 1 Rd
(13)
With = (k () 1) S. The relative position encodings are chosen as P; K.(1) = [ o]" for V(U)) 1 & i < k and p; K,01) 0 otherwise. Similarly, p; query matrices,
for every i. The hhth head has key and
W (1.))
0 Oix O
0
0
(14)
((LA)
(1x(k11)
0
0
With these choices, and letting - 00, the 1h layer computes the attention pattern,
att (1.A) lli n h).
(15)
n,
Choose the corresponding value matrix
(1.b)
O(2+n5)x2
0
(16)
0
Isxs
choosing the projection matrix appropriately, the output of the transformer after the first residual connection is,
x
0 01
1
exi
....
Czn- 1"
(17)
14

---

K.(2)
V.(2)
Layer 2. In this layer, the relative position encodings Pn- and query matrices are chosen as,
and Pn- are all set as O. The key
W Oskxk
(2)
1(SK+1)(((++1)))
0
0
0
0
(18)
W (k+x
(2)
[(S++1)((6+11)
0
0
With this choices, we have that,
W
(2) x
K Σ I( kn-s+1).
(19)
Taking x - :o, observe that the attention pattern concentrates on the uniform distribution over all coordinates i such that Zi-j In-s+11 for all j [A]] More formally, if Σi= [( xr,) > the attention pattern for any n > 1 is,
I[Vj E [k], zi-j = In-j+1)
2 R
(20)
E I(jj [k] Ti-j In-j+1)
The value matrix picks out ex from the embedding x; (Equation (17)) and places it in the last S coordinates. The subsequent linear layer picks out the last S coordinates, resulting in the logits,
I(Vj E [b], Fi-j In-j+1)
logit, Σ
(21)
VM Li I(Vj E [k]. Zi-, = In-s+1)
i=k+1
assuming that = T(Vý E [R]], Ii-j = In Øn-j+ j+1) > 0. In particular,
Σ"' Em=k+1 T(VO & j & ki In TT IT-
(ogit((DT+1)
(22)
2n 1(V1 < j & k. In-j IT-)+1)
assuming MT 1[V1 Ej k, In = 0T- > 0, i.e. the conditional k-gram model.
B Proof of Theorem 3
Define K* [1+0,((+++]]] by rounding &s 1 up the nearest power of and l log, (k*). In the setting of relative position encodings, given the sequence 21. In, while generating the output of the attention + feedforward layer for the symbol 2n. the embeddings æ; Emb(x,,, (,,), Pa-i are used for e [n]. 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 as,
x Emb(x,.) = 0 x
|
ex
01x((+ 0ixs I" (N"
(23)
where ex R is the standard basis vector in d' dimensions. And the relative position encoding for the keys as,
0JT
lixe=
if = 0, if (1,2,4,... kk* k2/2] otherwise.
P
(U)KK
(+log,(i)
o7
(24)
0ax11
And for the value vectors, p. = 0 for all i.
For the first layer and first head, will describe the value, key and query matrices. Choose,
w (1)
V
and,
K
0
(25)
(
VR
O
1 0
0
0
0
15

---

Then, observe that for i > 1,
(W (In-i + ( )..K). w xnn) xllle 1)
and for i = 0,
W (2) (1 Po ) w! xn
)n)) =
In particular, letting - co, the attention pattern is,
att (1) n. 2 2
1 I(i 0) I[ 1).
(26)
Choose the value matrix as,
s 0
21
together with the residual connection, we get,
2
0 01xe
S
xn
s 2n-1
0 (k* 2)S
01xs ]"
(27)
Layer l + 1. By induction, assume that the output of the thh transformer layer is of the form,
z (f+1) = 01xe
Vn
0 (x*
Oixs I"
(28)
for some vector U e g²'S We will show that with appropriately chosen key, query and value vectors the (E + 1)th layer, the output of this layer is,
(++2) 2)
01xe+
2n
Un + 2
0 xx 2
01xs ]"
(29)
We will consider the same relative position encodings and query matrix in this layer as in the first layer (Equations (24) and (25)). Consider key matrix of the form,
((+11)
O1xe
v
0
0
With this choice, observe that for
W (((+))
x n- +
(8+1)
+p, + (e+1). w x
K. I(i 22)
K
and for i 0,
W
(1(+))
(e+1) + Po (e+1)K w
x (e+1)) K
K
In particular, letting - co, the attention pattern is,
att (C+1) l +
n 2 2
I(i
29).
(30)
Choosing the value matrix as,
(=(+
0 0
0
21
we get,
2
(f+2)
O1xé+
1 Un
Un + 2n
0 01x(2=2241))5
01xs ]"
(31)
16

---

Final last transformer layer (L = E*).
The output of the second last transformer layer, indexed
1 is,
F/ x/
(2*)
(e")
0ixe
Un
(1-1)
Un
(8--1)
2 -1-
n 20-20-1
01xS
Oixe=
Un -1-
un
11"
(
01xs
which follows by plugging in the definition of k*. Note that there exists linear transformation (F)) such that.
2/2
11" ((() (2
x
01xe+
on
(1"
(1--1)
01xs
]"
This can be further decomposed as,
0ixe
Un
22)
(E*
Vn
22
+ 1
n-2
(6-2)
22
-2-
+ V n-
(6'
-2
0ixs
}r
And yet again there exists a linear transformation L [ee -1) which transforms this as,
(1"
L
zn
-1)
0 01xe
Un
-2)
v
2)
20
n x/- 22-
t1
(2)
n-
20-
01x8
]r
0 O1xe
Un
())
2)
-2) -))
22)
01x9
(32)
By recursing this argument and composing all the linear transformations, up to global permutation, we get that,
19 (e))
O Oixe
1
Un
nn
...
(1) (1
01xs 1"
T
01xs '"
T
O 01xe
1
(33)
In the final layer, we will right multiply the key, query and value matrices by L' = e= L The effect can be interpreted operating the original key, query and value matrices the embedding vectors in Equation (33). In the final layer, we will set all the position encodings to and consider the key and query matrices,
TT
(e') (e*)
-V
0 05xx((+++8))
Iskx Sk 0
O
(34)
(())
(C*)
Vk
O5Kx(e++1
Sk 0
0
Q
0
0
Then,
k- =
W
L'
('''''''''''' L' L'a((")
L'
(8) w
Σ
(35)
Where we must be careful note that the input æn contains copies of ex,, Cxn-1 Cra- since
k- > k + 1 by definition.
Letting 00, if there exists i such that En= T( = di---1) > 0. for 2 h. the attention pattern is,
att (1*)
1(x4-1 = Eni 24-2 = xn-1,
Ki-k = 2n-k+1)
(36)
n.2
[[ Ini 2 In 1,
Ti-k In-k+1)
17

---

Finally, choose,
W +2) ((-5)x((++1)
0
0
(37)
(1x()++1)
Isxs
0
get,
[[; In 1, -i kn-k+1) k+1)
0 (d-S)x1
x (f)+1)
(38)
2i=k [[ Bnz 21-2 zn ". Bi-k In-k+1)
ez.
1=k
Choosing the subsequent linear layer as,
A = (s [sxs]
(39)
b = 0sx1
(40)
Results in the output,
T - 2T-1; 2n-k 2T 2T-k-1) (ogit)(ET++) Σ
(41)
Ex 1(In-1 2T, 27-1,"" Ln-k IT++++)
n-
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 simple transformation which is to double the hidden dimension and add copy of -xn (e)
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,, 's is An O, and therefore the standard deviation the L2-norm of the features. order to avoid having to explicitly state this transformation each layer, we will simply redefine the layer norm LN to output e/ w/|lw] for the input vector v, which is realized on the first d coordinates of the transformed embeddings.
simply normalizes by
This transformation can be realized automatically by redefining the initial embeddings Emb(xn), and modifying the weights of the attention and feedforward subnetworks as follows: The input embeddings are changed to [Emb(x,,) -mmb(a,)]' e Iddd The key and query matrices are chosen to be 0 on the last d coordinates in every layer, the value matrix for > is transformed to Blkdiag({ ( w-(() }), and likewise changing the feedforward layer the block diagonal matrices Blkdiag( ( w/))) and Bikdiag( (! W)("). This transformation adds a copy of - (1) into the last d coordinates of the corresponding embeddings.
C.2 Notation and supplementary lemmas
For each i E [7]. define,
U: ExI-1 3 ex--2 + 3k- -x;
(42)
u: Cx; + +3 En; + .... + 34-1 Ex; k11
(43)
Note that although v; = u;-1, we make the distinction between the two to avoid any confusion in what is stored in the embedding vector at time and at time 1. Furthermore, define,
In = {k + +<+lliiin Vj [k], 21-1 = In-5+11.
(44)
Lemma 1. i In. 2; zn However, if k but $ In, then, |0:1
> 3 3-k
Let j E {0, 1......., 1} denote the largest index j such that Ān-s * I1-5-1- Consider the coordinates 2n-jr [S] and b zi-jt-1 E [S]. Then,
5
32
ea) (ui, Ex) 3
Σ
3
(45)
2
1=0
18

---

3
(2t,, eb) - (Ux, eb) 20
(46)
2
Unul2 2
U1l12. then,
U
Un
cb
(Ui, eb)
Vn Eb)
3
3
3
(47)
[1:
|
max uz
12
2
This uses the fact that u, and vn coordinate-wisg non-negative. On the other hand, if U 2 $ [|ulll, using a similar analysis,
Un
Ca
Eo
(48)
either case, there is coordinate (a or b) such that, u. u./l|u:lle and vn /n/l/ differ by at least 3 This implies the lower bound the L2 norm of the difference of the vectors.
C.3 Proof of Theorem 4
Choose the input embeddings as,
ar n
(1) Emb((r.) = (0 ege 01x5s] E TRSS3
(49)
In the first two layers will use the same relative position embeddings, in particular,
V Vlog(3)
if = 0, if E (1,2, ... k 13,
[1)KK pi
Pi
(2),K
(i+ 1) log 33
0
1
0
(50)
(k + 1) log (3)
0
if k.
and the value embeddings,
3'
[[
o]
DNV (2).V
Pi
(51)
0
i>
In the final layer, we will drop all position-relatee information and choose p (3) all i.
(3).V for
P:
Layer 1. Consider the key and query matrices,
w = Vn.
1ix2
0
0
(52)
((1) = vx.
O
1 x
0 0
0
0
Then, observe that, ((( (Emb(Fn-)) + P. ())K). w (Em((:.))) (m.))) = x(i + log(3) TTO & i min(n, k} - 1)
Letting c0, this results in the attention patter,,
3'T(0 & i min(n,k) 1)
att p,
(53)
Emin(n 3
Choose the value matrix as,
W (1) 0 ((+aa)x
0 1
The output of the attention layer (with the residual connection) is,
min [n.k)-
(
01x3
Un
01x35
|"
where, u,
Σ
att 6om..
(54)
1
19

---

In the feedforward layer to follow, we will choose,
(1) I
0 (3+25)
0
07
(55)
W (1)
=
Isxs
0
0_
Which simply extracts un from En With the subsequent layer norm and residual connection, the output of the first layer is,
2
0 01x3
S
Win u.
0 01x3s
|'
(56)
Layer 2. this layer, the relative position encodings and query matrix are the same as layer but the key matrix is chosen as,
0
1ix2 0
0
K
0
(57)
With this choice, observe that,
W K (2) P. (LJ.K)
x (2)
<(t + 1) log(3) I(1 & & R)
(58)
As before, since K o0, this results in the attention pattern,
3'1(11 & i min(k, n 1})
(2) att n.n-i
(59)
Ex=1 miin 3
which is similar, but subtly different from the attention pattern in the first layer (Equation (53)). The first layer focuses on indices i such that 0 & & ke 1, while this layer focuses 1 & & k. Choosing the value and projection matrices as,
I 3x3
0asx3
0
(60)
0
0
The output of the attention layer (with the first residual connection) is,
2 (2)
= Zn
0 xx
S zn
Un
Iu
My
Un
0 01x25
min min(k.nn1)) 1}
where, Un
Σ
attn,n-1 E
(61)
min(k,n-1)
and, Zn
Σ
att attz,n-i31
a short calculation to see that Z, 35+1/5 if 2 ++ and otherwise, Z, < 3k/5. This will be useful later, since the value of Zn, can be used determine whether n k n & which will allow the the next layer to avoid calculating the attention at & k. where the evaluation - is not well defined. In the subsequent FFN layer, we will choose,
Tn Tx-1," Ft1 k+1
w (2) I
0 (34 (3)+35)
0
(62)
W (2)
0
TsxS
0
0
0
0sx2s]
Which extracts Un from the embedding (2) connection, the output of this layer is,
With the layer norm and adding the final residual
T
(3) = Zn
02xl
Upy
0sx1
(63)
1
20

---

Layer 3.
In this layer, all the relative position encodings are set as 0 and instead,
V2k
0 (2-35)
Igx9
0
0
0
0
(64)
0
W
(3)
V2k=
0
5x x (2448)
Isx
0
0
0
With these choices,
2n<U;; U2 un
W
(3)
2/
(3)
W W (3) (3)
Q
||013] 2 | 12
Vi
Un
2xZiln + 2k - K
(65)
||1:112
2/,, 2
The resulting attention scores are,
(3) att n, oc exp
-Ki
Vi
Un
+ Zn
(66)
U2n |
Recall that Z, = (k + < n: Vj [k].
Øn-j+1 = xz-s}. Then for any i Ins v, = Un and
by Leemaa 1, for any i > 1 but not in In.
Un
14
3
Note that this gap small but non-zero. Furthermore, recall that Z, = 3 > k and otherwise 74 & 3* /5. Thus the attention prefers values of such that et; u, and such that k + 1. In particular, as - oo, the resulting attention pattern is,
Uni.(In)-
(67)
Choosing.
W
(3)
0
0
0
Osx3
Isxs
0
We get that,
(3)
att
(3)
0
S
(3)
0
I
The feedforward layer is chosen have (3))
(())
(3)
O, and the overall output of the final
transformer layer is therefore just (3) In the output linear layer, choose,
A Osx(d-S) Isxs]
(68)
b = 0
which results in,
[I Ej K Ii-s kn-j+t) Eni
logit
I
D =+1
[[V] $ k. Ze-j In-jt1)
sE
i=k+1
In particular,
DT 2n-kt1 T(VO & is k, In = xT-i+1)
(174a)
(69)
D T(V1 & & k. Fn-i = FT-++1)
which is the conditional k-gram.
21

---

D Representation lower bounds for 1-layer transformers: Proof of Theorem 5
We prove this lower bound by reduction communication complexity, and specifically to the set disjointness problem.
Suppose Alice and Bob given strings a. b 10, which are indicator vectors of sets A and B. Their goal is to jointly compute DIS( b) I(3i a; = b, 1), which indicates whether and B intersect not. Alice and Bob may send single bit message the other party over 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 this task.
Theorem 7 ((29)). Any deterministic protocol for computing DIS(a, b) requires at least rounds of communication.
We show that 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 a, b) in way which contradicts Yao's ioss lower bound in Theorem 7.
With T/3- suppose Alice and Bob have length m bit strings a, b e {0, 1)"". The transformer's input will be sequence of the form,
2, a,, b,, 2, 02 b2. 2, am, bm 2, 2,
(70)
of length 3m + = T 1. The input basically contains 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 Dr = 2. Noting that æT-1 and TT-2 1, the conditional 3- 3-gram computed to be,
[(. 1, = 1, 1, - 2)
(71)
1( 1,2,- 2)
Note that 2;-2 = 2, then must be the form 3j for j = 1,... .... n, and we may rewrite the sum as,
I[r3j = 1, 13 A
(72)
[1 1)
[B]
Now, let us use the transformer to construct deterministic communication protocol between Alice and
Bob, Alice is given T5, (33m-1) = (a), a2, amm) and Bob is given (I3, Te, .33m =
(b1, b2. bm).
In the first round, Alice computes the normalization in the softmax the attention which comes from the set of inputs she holds. For simplifying notation define.
score (h) (i) exp
(Emb(x;) + p:), w-(h) Emb(xr-1))
(73)
In particular, for each head h [H]. she computes,
((h)) log Σmn score (3) - 11)
(74)
Assuming that the transformer uses p bits of precision, Alice communicates ((h) 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,
(A))_ log (h)
2(h))
(h)
common
(75)
where 2(6h) = log Σ,- m and (h))
score (h)
(76)
log
score (3j - 2) + score (M) 2) + score score (T- -
m
(77)
which is the overall normalization term the softmax. This communicated back Alice, using another pll bits of communication. Next using this information, Alice computes the output of the
22

---

attention layer, taking the convex combination corresponding to the inputs she knows. In particular, for each h [H] she computes.
score (1) (3j - Emb(133-1) E IR..
(78)
exp( (1h)
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 V/lice which is d-dimensional. This is sent to Bob using dp bits communication. Subsequently, Bob computes the terms in the attention corresponding the inputs he knows as well the public inputs (all the 2's at positions 3j 2 as well as the last two symbols). In particular.
m score (h) (3j) (2a(( expl (())
m+11 score (b) (3j 2) Emb(2) score! (T (T-1) Emb(1) (AA(
(79)
exp (()))
These are yet again concatenated across all the heads and multiplied by the value and projection matrices to result in the output Y/ob which is added to VAlice to result in y. Bob passes 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 gram, which by Equation (72) equals |A n B). Note that if |A r B) B| is non-zero, it must be at least 1/".. This means, if the transformer is able to compute the conditional 3-gram to within an additive error of 1/37, then Bob can simply threshold the output of the transformer to decide whether A n B Ø not, thereby solving DiS(a,u) b).
Since this communication protocol is deterministic, by Yao's 's lower bound 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 2p/l dp + 1 (the last 1 comes from Bob having to communicate the answer to Alice), completing the proof.
E Lower bounds on representing h.th th -order induction heads: Proof of Theorem 6
In this section we prove the size-lower bound on attention-only transformers representing hthh -order induction heads in Theorem 6. To enable this result to better interpreted, we will break down into two corollaries.
Corollary 1. Consider an L--layer attention-only transformer with head per layer and relative position encodings, which satisfies Assumption 1. If + logg(k 22), the attention pattern in layer L of the transformer cannot represent a 1th -order induction head.
Corollary 2. Consider an 2-layer attention-onl only transformer with heads the first layer and relative position encodings, and assume that Assumption I is satisfied. If H & k 3, the attention pattern in the 2nd layer cannot represent a htth -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 -oh-rrder 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 He heads in layer & [L] 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-oni transformer cannot represent 1tth order induction heads for any l: > We will prove lower bounds on the transformer when the input is binary, i.e., S {0, 1). With relative position embeddings, observe that the first layer of the transformer model learns representations of the form,
(2) Emb(a,,) + Σ att (1)
Emb((:) + Σ W (1)
P
(80)
where note that the attention pattern only depends on n and i and not on 2; or Fn. These representa- tions are input into the second layer, which realizes the attention pattern att n, which is proportional
(2)
23

---

to,
exp ( (w 2) (z) +Pn-i
2)
k.(3),, W x2(2) (22)
(81)
Q
We need this function to be maximized uniquely when 2;-1 = En Zi = In-k+1. Denoting
(0)) w/ Emb(O) 0) and ((1) w Emb(1).
x Emb(xn)+ + Σ att w,
(1)
(1) Emb(
Σ w (1).V
Pn-i
(82)
= I, Emb( + (1 - nn ) Emb + Σ
att
n.t (z+ o(1) + (1 - (2)))(())))))))))))))))))))))()
+ Σ
W/
Pn-i
isn
(83)
Emb(1) (1) + Emb(O) ()
+ 2n
Emb (1) Emb(0)
+
Σ
att
att (1)
x(1) + 6(0)
+ x
x(1) - x(0)
2
2
2
2
+
Σ
w (1)
Pn-1
(84)
V
where x/ < 2x, 1. We can write this down as,
n
(2)
m + M [ar, 2n-1 x¹1²'
(85)
where _ is a matrix of rank at most 2 and of the form,
M (11) ( +10)
att n,1 + Emb(1) - Emb(O)
atto..
(1)
[1
0
(86)
2
which is independent of zq,," 5. Likewise m. (1) collects all the vectors in the sum that don't 't depend on z',-.. ... zy, Now, observe that in the next layer, we wish to show that an induction head cannot be realized by att for each i n. We will show this for any value of i k.
In the second layer, we may write down the key vectors as,
w (2) 2 + -.:))) w (2) m
(1) W (2) M, (1) [a]] 2/-1 211""
(2) (2),K W
Pa-i (87)
Again, defining the vector m!! (1) and the matrix M." appropriately (having rank most 2), this equals,
m (12 u ( x11 + +M (1)
(88)
where y [2/-1 " and the vector m. depends on z as well as the inputs T/-K-1.-... ,1, which in this context, are treated as nuisance variables since they do not inter- sect with 1z/-11''. , u (zn.... 111 }. Henceforth will avoid explicitly stating the (1)
dependency of m. on the zj's. Similarly, the query vector can be written down as,
w n = m + M x (
(89)
where (!! M and M are defined appropriately, with M and M. of rank at most 2, and is defined as [=',
" For appropriate matrix M n, vectors m and m n,5 and scalar m n.s+ m the dot-product of the key and query vectors can be written as,
W (2) x (2)
w (2) (2)
x M + T (M ) (mx,)" x + (mx + m & fn,( y),
(90)
Which is linear function in and quadratic in y. both of which lie on (+1)*. Note that the matrix
M i has rank at most 2 since it is product of M (1) and M. (1) each with rank at most 2. Next we introduce a lemma showing that M is inherently low rank, the quadratic form in Equation (90) which captures the dot-produc between the key and value vectors cannot satisfy the property that for every y, the function is uniquely maximized at æ = y. particular, this means that for any i E n - k there is some choice of Zn, "n- , -n k such that there are 1, -ii such that for at least one j E [k], Jız- and Tn-j-1 1 are not equal, but the attention score is larger than the case when zı-j were equal to -n-j-1 for each E [k].
24

---

Lemma 2. If M has rank & k - 2, it is impossible for (x. y) to satisfy the property that for every y E {+1}², the maximizer is uniquely
y.
The proof is almost complete: if k >> 4, then the rank of M n. 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 string such that the attention is longer uniquely maximized when -t-k-
In = 24-1," f+1
Proof. For the purpose of brevity, define Hx = {+1]%, First consider the reparameteriiation,
(M
T M where M R, = mm 7,
(91)
Then, the dot-product of the key and query matrices can be written as,
y 1 v" M (.,)) + (m,,,,) T y
m
(92)
Note that this function is linear in ã and therefore must be maximized on vertex of the convex hull of the domain, M Hr M,,,h h E Hz If M has rank at most k 2, the rank of M is at most A and cannot be full rank. We show that this must imply that there a vertex E Hx such that M n. v is not a unique vertex of the convex hull of Mr, This means that v cannot be a unique maximizer for ö when maximizing over all strings in Equation (92), and specifically y = e is witness Lemma 2.
Below we discuss how to find such a vector v. Note that M is not full rank, which implies that there exists a vector n such that M O. Without loss of generality, let n, be the smallest non-zero coordinate of n in absolute value. Then the vector n has no non-zero coordinates in the interval (-,,1). We will show that in n)) is good choice for v. Consider two cases,
Case I. Every non-zero coordinate of ni is in [+1]. Consider any i E Hr. which matches with on the non-zero coordinates. Consider x' which is the same æ. except negation is taken on the coordinates where is non-zero. Note that M v M a', for the same value of æc This means that for any y. In particular, from Equation (92), both æ and æ' 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, 'n n has non-zero coordinates which not all in [+1]. In particular, least one coordinate where this vector is strictly less than -1 or strictly greater than +1. In this case, observe that the sign vector ñ sign(r) mi n) He lies within, but is not vertex of the convex hull of the set He (n]'n). The reason for this is simple to see when we assume that n, ' n has only one coordinate which is not in [-1, 1]. say, the coordinate j = 2: here, ñ can be written down as convex combination (with non-zero coefficients) of n. 'n and (2); the latter vector is obtained by flipping coordinate 2 of î. When there is more than one coordinate not in [-1, 1], we can peel away these large coordinates in ni by taking a convex combination of this vector with the vectors ( for the appropriate values of j. to return the sign vector ň. Here, %0) is the version of ñ where the jth coordinate is flipped. This results in the following claim.
Claim 1. The sign vector lies within the convex hull of the points H4 U Uini" n),, but is not a vertex of this set.
In particular, we may write,
à = an n
Σ an
(93)
jE[n]
where ag > and Σド= a4 a2= By left-multiplying this on both sides by M and noting that n lies in the null-space of this matrix, we get,
M n
Σ
0j
M M,
(94)
je[n]
25

---

where note that Eje(m) aj is strictly less than 1, since ag > 0. We may write this vector as,
M M
ag0+
Σ
aj
M x nl)
jE[m]
ag
M R.
h+
a a,M n, n0)
(95)
2
hENL
]c[m]
Since ag O, this equation implies that the image of ñ under M itself falls within conv (MMMM,,, n, H.). but is itself a vertex of this set. This means that can never be maximizer of fo. (-: y, and in particular when y n. thereby proving Lemma 2.
(-: y) for any
E.2 L-layer attention-only transformers with 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 grows as function of the depth of the transformer. In the case the 2-layer transformer, we show that cannot exceed 2. With the addition of more layers, we show that it cannot exceed 2/--!
Recall from the notation in Equation (85) that the output of the first attention layer is,
m (1) + (1)
[ach
4'a-1
(96)
where M Rdx has rank at most 2. Let us rewrite this as,
(
M (1)
x11 1T
(97)
where M E RdXT is causally masked to be 0's when it operates for all indices i > Note that even with this causal masking, ! has rank at most 2, as discussed in Equation (85). By induction, assume that the output of the (e 1)th attention layer is of the form,
xn (E) m
(e-1) + M T1.T
(98)
where F1.T d [2'r 2T-11
111² Passing (0) through the 4th attention layer, we get,
((+1)
(()) Σ a
(f)
(6) w(()) (a! +
(6) (011))
(99)
ml (e-1) + M 1)
ar. Σ
att w (0 (8-1)
ni
m
Σ (t W
ni
M 1f-1)
BIT
Σ att
W (f) (f.V
Pn-
(100)
Define,
m!()
m
(t-1)
(f)
( m m((1))
(e-1)
Σ
att (
W(O
Pa-i and,
V0.V
(101)
ni
n,
M
(E-1)
Σ
att w ( ((((11))
(102)
ign
Then, we can write down,
(E+1) m(!))
M M (4)
N1.T
(103)
We also inductively assume that for every & n.
M. (E-1) has rank R 2f-1 and,
(ii) M 1) can be factorized in the form R ur of where only the v., depend on but
the 's do not depend on i.
26

---

Both of these conditions true when 1-1 1 as evidenced by the structure of M, in Equa- tion (86) and noting that M! (1) is obtained from M, by right multiplying by diagonal mask matrix. Using the recursion in Equation (101), we prove that the induction hypotheses (i) and (it) true at layer l as well. In particular using the decomposition in (ii), observe that,
M 99
Up
n,r
Σ 2n
att
ni
Σ r=1
up
Vir
(104)
R (t)
Σ
U4
W
Up
EEs
isn
att n., Vi.r
(105)
r=1
2R
Ur
(106)
where for r' E [R], UR+r
1 ur and Un.y" Lign att, n. Vi,-- Since M is the sum of 2R rank 1 matrices and therefore has rank most 2R & 2 proving both parts of the induction hypothesis.
By induction, at the end of the (L 1)th layer, we have an output which looks like,
a m (1-1) + M Z1.T
(107)
where M((-1) has rank at most 24-1, More importantly, note that by the causal masking, even though it appears depend the whole input sequence through F1.7, note that x.(") only depends on Ø1, ... Zn and not on the future inputs to this time n. In particular, by a similar argument as in the 2-layer case (ef. Equation (85) to Equation (90)), for any & n k dot-product of the key and query vectors at the th layer as a bilinear form which looks like,
decompose the
(W(l) (L) (x + + (L).K) ( x0
x TM
y MM,yy
las (x, y) (108) T" M yy + (m,,)) (m²)) Ty + mūs & p(l+t1)
where æ and y are defined as [a
d 1" and [2/-1 x//=]" ld respectively, and M has rank at most that of M (4-1) which 2 particular, 2f & k 2, by Lemma the proof concludes.
E.3 The general case: Transformers with H, heads in layer a Proof of Theorem 6
The hh head of the first layer of the attention-only transformer learns patterns of the form,
(1,A)
nn
Σ 1m
att (1,b)
w
(1,b)
Emb(x;; +
Σ
(1,A) V.(LA)
Pa-i
(109)
isn
( (0) + (1)
8² (1) -1))
Σ 2
(1.4)
x
Σ 18
w.l. V.a..) Pn-
(110)
2
2
where the last equation assumes binary input sequence, defines x 2z, 1 and uses the notation (a( = w(O. Emb(0) and (1) = (1.h) Emb(1). We can further rewrite this as,
x ((L,)) = m...))
(1,4) + (1.A) x1.T
(111)
where each (1.A) [gxx is rank and applies a causal mask on the inputs x; for > n. Recall that the output of the first attention layer applies a projection matrix the concatentation of (11A) across h e [H,] and then adds a residual connection. The output can be written down as,
(1.1)
M (1.1)
my
Emb Emb(x,, (())
+ w
EI.T
(112)
mn (1,
M (1,H)
27

---

m (2) + 1
(113)
where,
M
Emb(1) ))= + Emb(0)
en + W
and,
(114)
2
M
(1.
m (1.1)
m (1)
Emm(l) - Emb(O)
2 + w-(1)
(115)
(1.H1)
Notice that the rank of the matrix ( is most H1 + 1. This because the concatenation operation
can increase the rank at most additively, and since each of the M. (l.A) matrices are rank at most 1.
Following through the proof in Appendix E.2 for the L-layer case, we can prove inductively that any layer f, the output looks like,
(f) m + M WIT
(116)
where the rank of 0 is [H 1). Invoking Lemma 2, If + & k - 2, the attention-only transformer cannot realize -oddr induction head layer L.
F
Model architecture and hyper-parameters
The experiments were run on one 8 x A100 GPU node.
Parameter
Matrix shape
wra
d x
transformer.wpe
transfomerr..n.. (xx)
d x 1
tranformer.battm.... (x
3d x d d d d
munsformer.batttn...ooo (xé)
transformer...n. (x
tansformer.h.mpp... ()
4d x d d 4d d dxl
mansformerh.mmlp...o (x()
transformer.nn
Table 2: Parameters in the transformer architecture with their shape.
G
Additional experimental results
Assumption suggests that the attention patterns att in layers ( 1,2,.... L- 1, as learnt by an L-layer attention-only transformers may only a function of only the position indices n, i. this section run additional experiments to test this conjecture. We train layer attention-only transformer with i heads in the first layer, on data drawn from a randomly sampled -od-order Markov process, and focus the learnt attention patterns plots the results of this experiment for h: and Figure 8 for l = 3. While in both cases there is some variance in the attention patterns learnt by the transformer in some of the heads, believe that this is a consequence of the iteration budget of the transformer, and specifically the fact that even if the test loss appears 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 function of the input, relaxation of Assumption 1, also believe that the results proved in Corollaries and 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.
function of in the input sequence. Figure 7
28

---

Dataset Architecture
k-th order binary Markov source
Based the GPT-2 architecture implemented in [30]
Batch size Accumulation steps
Grid-seane sradchearched in {8, 16)
Optimizer
AdamW (31 0.9, 82 0.95)
Learning rate
0.001
Scheduler # Iterations
Cosine Up to 25000
Weight decay
x 10-3
Dropout
Sequence length
Grid-searched [32, 64, 128, 256, 512, 1024) Grid-searched {16, 32, 64)
Embedding dimension
Transformer layers
Between 1 and 8
Attention heads
Upto to k
Repetitions
3
Table Settings and parameters for the transformer model used in the experiments.
0.35
0.25
0.30
Alcstion wode
0.15
0.10
0.10
0.05
10
Sequence index
Sequence index
(a) First head
(b) Second head
Figure 7: Mean attention for column n = of the two heads of the first attention layer, for a 2-layer 2-head transformer model trained on order-3 Markov process, averaged across 100 input sequences of length 128.
0244
632
0.31
615-
2.50
0.24
0.28
0.20
0.54
6.55
605
632
Soquenee indox
Sequence index
Sequence indes
(a) First head
(b) Second head
(c) Third head
Figure 8: Mean attention for column 10 of the three heads of the first attention layer, for 2-layer 3-head transformer model trained on an order-3 Markov process, averaged across 100 input sequences of length 128.
29

---

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]
The paper presents experimental evidence about ht -oder Markov processes being effecttiely learnable by small-depth transformers. We also prove theorems about the representation power of low-depth transformers, which stated in the main paper.
Justification:
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We state that our lower bound in Corollary is a conditional result in Sec- tion 6.1, and that 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 dynamies of gradient descent 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 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 vell-specifiiattio, 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 few datasets with 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
30

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
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 include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch 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 properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Experiment hyperparameters are presented in Tables 2 and and the code has been provided along with the supplementary material.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, a No answer to this question will not 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 dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
31

---

reproducible, they In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the code has been open sourced; hyperparameter choices have been provided in Tables and 3.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips. //npss.ccppubltii. wwwpcc/cpcpppppeeecccecccocceeep http://///didee//eeeubibiilt /oodeeooboisssssseeooee ssionPolicy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Tables 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 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 plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
32

---

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.
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 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 range (e.g. negative еπτог 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 text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details in Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, 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 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 the paper conform, in every respect, with the NeurlPS Code of Ethics https httpa::///////pe/eeeeeeee cctppcccccc//e/eetccccccccclcllt http:///thhhicieeenneeeeeenneeee
Answer: [Yes]
Justification: No NeurIPS code of ethics were violated.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations 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 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.
33

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No models with high risk for misuse were trained or released.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include URI.
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.
34

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datastts has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset 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. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included 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 equivalent approval/review based the requirements of your country 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 research with human subjects.
Depending 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.
35

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
36


Reviews:
[{"review_id": "OftiWnf20j", "reviewer": "Reviewer_DmcT", "summary": "The 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.", "questions": "1. 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?", "limitations": "Yes, the authors have discussed the limitations separately in the Appendix.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Given 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and questions about the paper. Below we address the main questions and weaknesses raised:\n\n### **[W1] Comparing single vs. multi-head**\n\nIn the attached material, please refer to Fig. 1, which we will add into the paper in the subsequent version (as the new Fig. 1). This plot discusses how the multi-head transformer learns, while the number of iterations to converge to optimality is of the same order, the test-loss these models achieves is nearly $0$. In comparison, the single head $2$-layer model has non-zero but small test-loss, while the single head $3$-layer model has even smaller test loss. At this scale, we did not distinguish between the small loss of the transformer vs. negligible loss incurred by the multi-head setup. It is indeed plausible that with better choices of hyperparameters (optimizing over embedding dimension, for instance), we may see the delta become even smaller.\n\nIn our evaluation setup, the transformer does not know what the distribution of samples at test-time are (only that they are Markovian). Thus, any model which achieves low test-loss *has to* resort to something like an in-context estimate. We were computationally bottlenecked to train transformers on sequence lengths of up to around $500 \\approx 2^9$ for the number of iterations required for the test-loss to begin to converge. At this scale, the model has “just” enough data to be approximately able to estimate order-8 Markov chains in context (i.e., having to approximately estimate $2^8 \\approx 250$ conditional distributions on $\\{ 0,1\\}$). At this scale, if were to increase $k$ the model would stop performing well, and more importantly, no in-context estimator could work at this scale. Information theoretically, there is just not enough data in a test-sequence to be able to estimate the next sample distribution, even approximately.\n\n**TLDR;** When the sequence length is around $500$, $k=8$ is the point where we would start to see degrading performance (for any model, and not just a transformer) by virtue of approaching the regime where we don’t have enough samples to estimate conditional probabilities from the test-sequence accurately.\n\n### **[W2] Long training to observe contradictory results**\n\nThis is discussed in the common rebuttal above.\n\n___\n\n### **[Q1] Why do 2-layer 1-head transformers learn order-4 Markov chains?**\n\nThis is a great question. While our construction shows that the standard $3$-layer transformers with $1$ head are able to represent $k$-th order induction heads, it may be possible that at an even smaller scale (say $2$-layer with $1$ head) that transformers are still able to represent the in-context conditional $k$-gram with some amount of error. It is also possible that at an even smaller scale, (say $2$-layer $1$-head transformers for sufficiently large $k$, or if the embedding dimension is reduced even further) that induction heads are no longer representable at all. Understanding this regime presents new avenues of research: here the transformer may still achieve low loss, while “provably” not being able to use the induction head mechanism to achieve this loss.\nIt’s worth pointing out the case of $1$-layer $1$-head transformers which were studied in [1]. The authors show that transformers are still able to learn Markov chains, but under this extreme size constraint, learning does not happen in a transductive way. At this scale, the transformer does not do in-context learning, rather learns the parameters of the Markov chain from which data is generated directly. This is a slight departure from our setting, since learning the parameters is infeasible when the training and test data come from different distributions (such as two different randomly sampled Markov chains).\n\n### **[Q2] Precise dependency on embedding dimension**\n\nThis is discussed in the common rebuttal above.\n\n&nbsp;\n___\n### References\n[1] Makkuva, Ashok Vardhan, et al. \"Attention with markov: A framework for principled analysis of transformers via markov chains.\" arXiv preprint arXiv:2402.04161 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JVjpYQg7la", "reviewer": "Reviewer_YGsY", "summary": "This 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.", "questions": "1. 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?", "limitations": "The 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.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. 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)", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments. Below we address the main questions and weaknesses:\n\n### **[W1] Error bounds**\n\nWe are happy to include a longer discussion of error bounds in the paper. To expand, suppose all the weights in the transformer model are upper bounded by $1$. When the transformer is implemented with a precision of $m$ bits, the new attention weights $\\widehat{\\text{att}}_{n,i}$ satisfy,\n\n$$\n\\frac{1+e^{-\\varepsilon}}{1+e^{\\varepsilon}} \\le \\frac{\\widehat{\\text{att}}\\_{n,i}}{\\text{att}_{n,i}}  \\le \\frac{1+e^{\\varepsilon}}{1+e^{-\\varepsilon}}; \\quad \\text{where, } \\varepsilon \\triangleq d \\cdot 2^{-m+1}\n$$\n\nThis uses the fact that by truncating vectors to $m$ bits, the approximation error in the inner product of the inner product of two $d$-dimensional vectors. i.e. $\\varepsilon$, is at most $d \\cdot 2^{-m+1}$. In particular, simplifying assuming $\\varepsilon$ is small, we get a multiplicative error of $[ 1 - \\Omega (\\varepsilon), 1 + \\Omega (\\varepsilon)]$. A similar analysis works for layer-normalization and results in a similar error scaling. Thus, within a single attention layer, we compute vectors which are within $[1-\\varepsilon', 1+\\varepsilon']$ where $\\varepsilon' = T \\cdot 2^{-m-\\log(d)}$. Iterating across $L$ layers, the error incurred scales as a multiplicative $1 \\pm 2^{-m-L\\log(dT)}$. Thus, when the bit-complexity scales as $m \\gg L \\log (dT)$, the approximation error is a multiplicative polynomially small constant in $T$.\n\nThere is only one caveat in case of the constant depth transformer, the weights of the transformer in each of the $3$ layers are upper bounded by $O(2^{k})$, rather than $1$. Thus, the bit precision must now scale as $k + O(\\log(ST))$.\n\nWe chose to de-emphasize these details in the paper to avoid unnecessarily complicated theorem statements, and to avoid detracting from the (arguably more interesting) other phenomena discussed in the paper.\n\n### **[W2] Intuition behind Theorem 5**\n\nWe have rewritten this section to make the intuitions more clear. We have added a new “proof sketch” section, clarified what the $k$-th order induction head does in this context, and emphasized how layernorm allows the transformer realize these $k$-th order induction heads.\n\nIn short, layernorm allows changing attention from $\\text{att}\\_{n,i} \\propto \\exp (\\langle k_i, q_n \\rangle)$ to instead look like $\\text{att}\\_{n,i} \\propto \\exp ( - \\| \\hat{k}\\_i - \\hat{q}\\_n \\|\\_2^2)$ where $\\hat{k}\\_i$ and $\\hat{q}\\_n$ are the $L_2$-normalized key and query vectors. This is nice, because if we can realize the key vectors $k_i = \\sum_{j=1}^k 2^j e_{x_{i-j}}$ and the query vectors as $q_n = \\sum_{j=0}^{k-1} 2^j e_{x_{n-j}}$, the attention is maximized if and only if $\\hat{k}_i = \\hat{q}_n$, which we can prove occurs only when $x\\_{i-1} = x\\_n, \\cdots, x\\_{i-k} = x\\_{n-k+1}$ (using the fact that the binary representation of a number is unique). This realizes a $k$-th order induction head.\n\n### **[W3] Representation vs Training dynamics**\n\nIn this paper, we focus on the representation capacity, which itself turns out to be a complicated phenomenon. On the optimization side, there are indeed many questions open. Unfortunately, even in the simplest settings (1 layer transformers), the analysis of gradient descent on transformers trained on Markovian data is incredibly complex. The rigorous analysis in [Makkuva et al] is over 70 pages long for this case. Extending to higher depth, is incredibly non-trivial. [Lee et al] discusses this case, however, their analysis is only amenable under very strong assumptions i.e., the “reduced model” of transformers. Even with these assumptions, the analysis is over 60 pages long. We hope to convince the reviewer that while training dynamics is indeed an interesting and timely question, it certainly requires a dedicated effort to be able to analyze, and would constitute a separate work of its own. \n\n### **[W4] Experiments on higher values of $k$**\n\nDiscussed in the common rebuttal.\n\n___\n\n### **[Q1] Learning dynamics**\n\nAt a high level, yes, we believe that the k-head mechanism discussed in prior works may be more amenable for analysis. When the depth is constant, the transformer is forced to be more creative in the ways it learns the $k$-th order chain, and this leads to more intricate mechanisms. However, unless simplifying assumptions are imposed, we believe that understanding learning dynamics falls into the realm where current tools in optimization theory are not strong enough to be able to answer.\n\n### **[Q2] $\\Omega(k)$ bit precision for learning $k$-th order Markov processes**\n\nOne may consider an extension of $k$-th order Markov processes to the case where the conditional distribution depends on a sparse subset (of size $p$) of the previous symbols observed. The conditional distribution has a low degree of dependency on the past, even though it is not the immediately preceding $p$ symbols. While these kinds of processes are special case of $k$-th order Markov processes, modeling them as such may necessitate a very large value of $k \\gg p$, since the there may be dependencies on symbols which appeared long ago.\n\nOur constant depth constructions work even in this setting, with the bit-complexity scaling with $p$ and not with $k \\gg p$. This analysis shows that transformers don’t require high bit-precision as long as the degree of dependency of the conditional distribution on the past is small. In practice, this is often the case.\n\nOn a separate note, our constant depth constructions require $\\Omega(k)$ bit complexity, but the embeddings are now $O(S)$ dimensional, rather than being $\\Theta(Sk)$ dimensional in the attention-only setting. Thus, the total number of bits in each embedding vector is the same in both cases. We believe that it is possible to reduce the bit complexity of the constant depth construction at the cost of making the embeddings $\\Theta (Sk)$ dimensional.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EuEJlcoOWr", "reviewer": "Reviewer_873s", "summary": "This 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.", "questions": "- 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.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed comments.\n\n### **[W1] Thorough discussion for why contradictions arise in past work**\n\nAs discussed in the common rebuttal, the main reason for the differences in observations compared to [1] and [2] is the fact that the models were not trained for sufficiently long to observe the increase in test loss. Please refer to Fig. 2 in the attached document where we plot this explicitly on a random $5$-state order-$3$ Markov chain. In the beginning phase of the training (iterations 1-200), the test-loss curves plateau around the bigram performance, but as we continue training, the loss continues decaying slowly. We believe that this difference only appears because the model was not trained for sufficiently long, keeping all other variables unchanged. We will include this plot in the paper.\n\n___\n\n### **[Q1] Range of $k$ considered in previous work**\n\nThe contradiction we mention in the $2$-layer case is empirical. Prior work such as [1],[2] train $2$-layer $k$-head models on $k$-th order Markov chains for $k=1,2,3,4$. However, we are lead to believe that the $2$-layer $1$-head model was not trained for sufficiently many iterations to learn $k=2,3,4$.\n\nFig. 2 in the attached material showcases this difference when a $2$-layer $1$-head transformer is trained on $k$-th order Markov chains on $5$ states for $k=3$. The initial phase (first ~100 iterations), the model stagnates in performance at the level of the best bigram model. When we continue training, the loss of the model starts decreasing further and breaks past this bigram phase to learn higher order conditional $k$-gram models.\n\nFinally, it’s worth noting: while we do not have a column for the $2$ layer standard transformer in Table $1$, we believe that the main $3$-layer construction can also work for $2$-layer transformers with $2$ heads in the first layer. This is because in the $3$-layer construction we consider, the first two layers of the model essentially do not interact with each other (and operate on different parts of the embedding space), they can be implemented across two different heads of the first layer. We will add this into the paper to clarify that there exists a contradiction even theoretically.\n\n**Formally:** For $2$-layer transformers with $2$ heads, prior work argues that these models can learn up to $k \\le 2$. Our work shows that these models can represent much higher values of $k$ theoretically (via an extension of Theorem $4$).\n\n### **[Q2] Clarifications regarding Fig. 2**\n\nWe have included the plot for $k=1$ in Fig. 1 of the attached document in the paper (this experiment trains $k$-th order Markov chains on $k$-head $2$-layer transformers). We will include this plot in the paper as well. There is a gap in the test loss which decreases when we move to $3$-layers and $1$-head. The reasons for this small but non-negligible test loss gap are unclear - while our theoretical construction works for the $3$ layer case, it may be possible that with $2$ layers and $1$ head transformers are unable to exactly represent the conditional $k$-gram model. It may also be possible that ADAM/GD are unable to find this solution, even if it is exactly representable. Understanding the limits of how $2$-layer $1$-head models learn is an interesting question to look into.\n\nOn a separate note, It is worth mentioning that the best achievable test loss will increase as we keep the sequence length fixed and increase $k$. In the learning setup we consider, the training and test data do not come from the same Markov process, since the transformer is trained on data from randomly chosen Markov processes and tested on data from another randomly chosen Markov process. So the transformer *has* to do some form of in-context learning. When the sequence length is fixed and $k$ increases, the in-context learning problem becomes harder. The model has to learn $S^k$ conditional distributions (for each possible prefix of $k$ states). So as $k$ grows, the model has less data for each of these distributions which makes the estimation task harder. This is a major reason for why the $3$-layer transformer’s loss increases a little bit when dealing with $k=8$. We were computationally bottlenecked to train on sequence lengths of up to $500$ and at this scale, the model has just about enough data to learn each of the $256$ conditional distributions. Fixing the sequence length as $500$, the models should not be expected to have small test loss as $k$ grows to $9$ and beyond.\n\nThe reason we do not train on $k > 8$ is discussed in the common rebuttal. The sequence length to consider would have to be at least around $1000$ or longer for $k=9$ and $2000$ for $k=10$. The model also takes longer to train, and we were unable to train for long enough for the model to reach convergence. In the future, we hope to be able to test out the results in this paper on longer sequence lengths (thereby allowing larger values of $k$) while keeping the depth of the model and the number of heads fixed.\n\n### **[Q3] Additional suggestions**\n\nWe shall plan to incorporate the notational changes and fix the typos as suggested by the reviewer. With the additional page in the final version, we shall also include a discussion on related work on the role of softmax attention and depth which were present in existing work. Finally, we will also include a discussion on related works which consider techniques similar to the composition technique we considered in Theorem $3$.\n\n&nbsp;\n\n___\n### References\n\n[1] Makkuva, Ashok Vardhan, et al. \"Attention with markov: A framework for principled analysis of transformers via markov chains.\" arXiv preprint arXiv:2402.04161.\n\n[2] Nichani, E., Damian, A. and Lee, J.D., 2024. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735.\n\n[3] Edelman, Benjamin L., et al. \"The evolution of statistical induction heads: In-context learning markov chains.\" arXiv preprint arXiv:2402.11004.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jZb1DlMtAY", "reviewer": "Reviewer_uqDv", "summary": "This 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).", "questions": "1. 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?", "limitations": "See weakness.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments, and constructive criticism. Below we address the main weakness pointed out, as well as the question asking about experiments for $k > 8$.\n\n### **[W1] Evidence that ADAM/GD converges to theoretical construction**\n\nThis is a great question. In Fig. 3 of the rebuttal pdf attached, we plot the attention patterns learnt and observe some sort of approximately exponential decay. The theoretical constructions also match this kind of exponential decay (as seen from the attention pattern in eq. (50) in the proof of Theorem 4). While this is promising, there is certainly more to explore about how and what the transformer learns. In particular, it is interesting to see how these attention patterns evolve over the course of learning before converging at these kinds of exponential patterns. This may reveal some more about how transformers learn, and get us one step closer to a theoretical understanding of the learning dynamics of transformers when exposed to Markovian data.\n\n### **[Q1] Experiments for $k > 8$?**\n\nThis was also pointed out by other reviewers, and is indeed an important point. We would be happy to include something which provides more clarification around what happens when $k=8$, however in the form we currently consider, it is not within our computational budget to scale this to $k=9,10$ or anything higher. A longer discussion as to why this kind of computational “phase transition” occurs is provided in the common rebuttal above.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6S5gqidQjk", "reviewer": "Reviewer_TFSx", "summary": "This 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.", "questions": "(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?", "limitations": "(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.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(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.", "weaknesses": "(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)$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and constructive criticism. We have fixed the typo pointed out. Below we address the main weaknesses discussed in the review:\n\n### **[W1,W2] Experiments for and learning dynamics under larger state spaces?**\n\nIn the attached rebuttal pdf, we ran experiments for Markov models on state spaces of size $|S| = 5$. Here too, we show that the main phenomenon discussed in the paper ($2/3$-layer $1$-head transformers learn $k$-th order Markov processes for $k > 1$) holds true. In particular, we show that a $2$-layer $1$-head transformer learns order-$3$ Markov processes. We were unable to run this experiment for $k=4$ keeping the state space size fixed since the model longer to converge than our computational allowance would permit. When the state space becomes even larger, say on the order of tens or a couple of hundreds, it becomes much harder to train the models on reasonably large values of $k$ within the computational budget in our availability.\n\n___\n\n### **[Q1] Binary state space**\n\nYou are correct to note that we run our experiments on binary data (to be able to scale to the largest possible values of $k$ on a computational budget).\n\n### **[Q2] Precise dependency of embedding dimension**\n\nThe experiments we run are more in line with the case of the standard transformer architecture, which corresponds to Theorem 4, where the embedding dimension scales as $O(S)$ (where the constant is implicit). In the common rebuttal, we discuss how to theoretically reduce the dependency on the embedding dimension at the cost of a small increase in the loss incurred by the model.\n\nOn a separate note, there are usually a few reasons to expect that there would be some differences in theoretical and empirical findings especially when it comes to the precise dependency on quantities like the embedding dimension, etc. Empirically, transformers are trained from a random initialization and in a way where the model may be more / less efficient at realizing certain mechanisms compared to a theoretical construction. We would argue that it is in fact promising that the embedding dimension is of a similar order as what is predicted in our theorems, as it is often the case that the theoretical and empirical findings are off by several orders of magnitude.\n\n### **[L1] Time inhomogeneous processes**\n\nThis is a great question, and something we hope to study as the next step. Time homogeneity of the data model (such as with $k$-th order Markov processes) implies that it is possible to write down the estimate of the next-symbol conditional distribution using an unweighted empirical estimate. When the data process becomes time inhomogeneous, the statistical estimators would look more akin to weighted conditional empirical estimates (under assumptions on the nature of time inhomogeneity). The constructions we consider in the paper do have the ability to translate to these kinds of settings as well. However, to understand the limits of how transformers do this is a great direction for future research, and something we hope to explore going forward. It is plausible that transformers need more depth to be able to capture these kinds of processes, depending on the nature of the time inhomogeneity.\n\n### **[L2] Constructions only apply for small state spaces**\n\nIn commercial (for which metadata is available) and larger open source models, the vocabulary size has usually been in the range 30K-250K. However, the embedding dimensions have been of the scale of 1000-20,000, which is one order of magnitude smaller. It is worth pointing out that our constructions, while optimized for the dependency on the depth and the number of heads, can also be improved on the depth dependency. In the common rebuttal, we discuss how to improve the dependency of the embedding dimension from $O(S)$ to $O(\\log (S)/\\epsilon^2)$ where $\\epsilon$ is a parameter which now captures an approximation error. Thus, even when the state space is larger, for practical ranges of $\\epsilon$, the dimensionality implied by this result should be much smaller than $O(S)$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## **Common Rebuttal**

We thank all the reviewers for taking the time to go through our paper and suggest constructive criticism. Please find attached a pdf containing additional plots.

Below we address some common points raised by multiple reviewers.

### **Long training reveals contradictions with prior work**

We point the reviewers to Fig. 2 in the rebuttal pdf. Here, we clearly observe the stages of learning for a $2$-layer transformer. Initially (in the first ~200 iterations) the model stagnates and this is around the loss of the best bigram model (which is what previous work claimed these models could learn). As we continue training, we see that the loss starts decaying but over a much slower time-scale. The model indeed requires around 20K iterations to get close to the optimal loss. The closest related phenomenon to why the model performance appears to improves very slowly is that of grokking: zooming into any interval, it may appear that the model has converged, however this only appears to be the case as the parameters of the model may be in a region where the loss landscape is very flat. 

In the rebuttal pdf, we also point the reviewers to Fig. 1. The authors of [1] also present similar results, but it appears that their experiments requires more samples / optimization cycles than in our paper. Here, we observe that $4$-head transformers with $2$ layers need around 3K iterations to learn $k$-th order Markov chains for $k=4$. In contrast, according to Fig. 2 in the main paper, $2$-layer transformers with $1$ head trained on the same data processes require around 7K iterations.

### **Running our experiments on higher values of $k$**

Running our experiments in Fig. 2 on higher values of $k$ is certainly important. However, there is a computational tradeoff which we were not able to navigate when $k$ grows to be larger than $8$. We were computationally bottlenecked to train transformers on sequence lengths of up to around $500 \approx 2^9$ for the number of iterations required for the test-loss to begin to converge. At this scale, the model has “just” enough data to be approximately able to estimate order-8 Markov chains in-context. Note that the model has to approximately estimate $2^8 \approx 250$ conditional distributions to learn order-8 Markov chains. At this scale, if were to increase $k$, keeping the sequence length fixed, the model would stop performing well. This is the scale at which no in-context estimator could work. Information theoretically, there is just not enough data in a test-sequence to be able to estimate the next sample distribution, even approximately. In order to test our results, say for $k=10$, we would have to evaluate on transformers of sequence length $\approx 2000$ (4x blowup). While the forward/backward passes naively takes 16x longer, the transformer also seems to require many more iterations to optimize the loss to low error. These reasons make it prohibitive to evaluate on larger values of $k$.

### **Precise dependency on embedding dimension, especially when $S$ is large**

In practice, the embedding dimension of transformer models like GPT-2/3 and LLama-2/3, and plenty of other open source models (for which data is available), the embedding dimension falls roughly in the range (1000-20,000). This is an order of magnitude smaller than the vocabulary size of these models, which is often in the range (30,000-200,000). It should be possible to improve our results in the dependency on the embedding dimension, which we didn’t fully optimize for (in comparison with the number of heads and depth). At a high level, the reason the embedding dimension in our Theorems  and ___ scales as $O(S)$ is because the transformer stores information about the statistics of each symbol along orthogonal components of the embedding vector. The orthogonality of these components prevents information about the statistics of any one symbol from affecting those of another symbol.

It remains to be formally verified, but we believe that all of our constructions, as well as the mechanisms empirically learnt by transformers are robust to errors which may appear even when information about the statistics of each symbol does not appear in exactly orthogonal components of the embedding vector, but in *approximately* orthogonal components. By allowing for approximate orthogonality, the transformer can be much more efficient about its use of the embedding space: the Johnson Lindenstrauss theorem states that there are $S$ approximately orthogonal (pairwise inner products $\in [-\epsilon,\epsilon]$) vectors in $\mathbb{R}^p$ even when $p = O(\log (S)/\epsilon^2)$ is much smaller than $S$. In principle, it suffices to choose the embedding dimensions to scale roughly as $\log (S)/\epsilon^2$ to have an approximate version of Theorem $4$, paying a error scaling with $\epsilon$.

Finally, it’s worth mentioning that for practical ranges of $\epsilon$ and $S$, this number ($\log(S)/\epsilon^2$) can indeed be much smaller than what is predicted by what is in the current version of Theorem $4$.

We are happy to include this discussion in the paper to provide more context about the embedding dimension.

&nbsp;

___

### References
[1] Edelman, Benjamin L., et al. "The evolution of statistical induction heads: In-context learning markov chains." arXiv preprint arXiv:2402.11004.

Meta-review: The paper shows expressibility and empirical results of Transformers when the input context follows k-th order Markov chains (i.e., the in-context learning setting), and characterizes the number of layers required to learn high-order chains from the theoretical point of view, showing surprising capability of Transformer of log_2(k) layer to encode k-th order Markov chains. Reviewers all agree the direction addressed by the work is important and novel.
Decision: Accept (poster)