What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains
Authors: Chanakya Ekbote
Ashok Vardhan Makkuva
Marco Bondaschi
Nived Rajaraman
Michael Gastpar
Jason D. Lee
Paul Pu Liang
Venue: NeurIPS 2025 spotlight ()
Keywords: transformers
interpretability
Primary area: deep_learning
TL;DR: The paper proves that a two-layer, single-head transformer can reliably perform in-context learning on any-order Markov chains.

Abstract: In-context learning (ICL) is a hallmark capability of transformers, through which trained models learn to adapt to new tasks by leveraging information from the input context.  Prior work has shown that ICL emerges in transformers due to the presence of special circuits called induction heads. Given the equivalence between induction heads and conditional $k$-grams, a recent line of work modeling sequential inputs as Markov processes has revealed the fundamental impact of model depth on its ICL capabilities: while a two-layer transformer can efficiently represent a conditional $1$-gram model, its single-layer counterpart cannot solve the task unless it is exponentially large. However, for higher order Markov sources, the best known constructions require at least three layers (each with a single attention head) - leaving open the question: *can a two-layer single-head transformer represent any $k^{\text{th}}$-order Markov process?* In this paper, we precisely address this and theoretically show that a two-layer transformer with one head per layer can indeed represent any conditional $k$-gram. Thus, our result provides the tightest known characterization of the interplay between transformer depth and Markov order for ICL. Building on this, we further analyze the learning dynamics of our two-layer construction, focusing on a simplified variant for first-order Markov chains, illustrating how effective in-context representations emerge during training. Together, these results deepen our current understanding of transformer-based ICL and illustrate how even shallow architectures can surprisingly exhibit strong ICL capabilities on structured sequence modeling tasks.

What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains
Chanakya Ekbote* * I
EPFL
Marco Bondaschi
cekbote@mit. edu
Nived Rajaraman B
EPFL
Jason D. Lee
Michael Gastpar
EPFL
UliT
Ashok Vardhan Makkuva'
Paul Pu Liang?
B University of California, Berkeley
Princeton University
Uii Massachusetts Institute Technology
EPFL -FoLe Polytechnique Fédéralc de Lausanne
Abstract
In-context learning (ICL) is a hallmark capability of transformers, through which trained models learn adapt to new tasks by leveraging information from the input context. Prior work has shown that ICL emerges in transformers due to the presence of special circuits called induction heads. Given the equivalence between induction heads and conditional k-grams, recent line of work modeling sequential inputs as Markov processes has revealed the fundamental impact model depth on its ICL capabilities: while a two-layer transformer can efficiently represent a conditional 1 1-gram model, its single-layer counterpart cannot solve the task unless it is exponentially large. However, for higher order Markov sources, the best known constructions require at least three layers (each with a single attention head) leaving open the question: can two-layer single-head transformer represent any tt -order Markov process? In this paper, we preciscly address this and theoretically show that two-layer transformer with one head per layer can indeed represent any conditional k-gram. Thus, our result provides the tightest known characterization of the interplay between transformer depth and Markov order for ICL. Building on this, we further analyze the learning dynamics of our two-layer construction, focusing simplified variant for first-order Markov chains, illustrating how effective in-context representations emerge during training. Together, these results deepen current understanding of transformer-based ICL and illustrate how shallow architectures can surprisingly exhibit strong ICL capabilities on structured sequence modeling tasks. Code available at the ! link.
Introduction
"A complex system that works is invariably found to have ovovedd simple system that worked." John Gall, Systemantics (1975)
Transformers, powered by the attention mechanism, have emerged as the dominant architecture in machine learning, achieving state-of-the-art performance across wide range of domains, including
*Corresponding author.
EEuual contribution/ equal mentorship. Authors listed alphabetically; either may be considered last author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

natural language processing [6], computer vision [32], and complex reasoning tasks [23, 36]. A key factor underpinning this success is their ability efficiently model sequences and perform in-context learning (ICL))adee adapting to unseen tasks during inference by leveraging the relevant input context [22]. well known that ICL emerges in transformers due to the presence of special circuits called induction heads [22, 31].
(ICtinadapttng
Intuitively, these circuits enable transformers implement "copy-and- mechanism by copying earlier tokens from the input and matching them with the desired context for next-token prediction. For example, first-order induction head mimics the functionality [.....A, A, ........................ B. Capitalizing on the connection between induction heads and conditional k-gram models, a recent active line of work has leveraged -rder Markov chains as a simple yet powerful framework to analyze how transformers learn induction heads, with the framework referred to as Markov-ICL [10, 20, 26, 7, 5]. In particular, using Markovian input sequences, the goal is to understand how transformers learn to represent conditional k-grams-equivalent to -ordee induction heads (Def. 2).
Building upon the Markov-ICL setting, Bietti al. [5], Edelman al. [10], and Nichani et [20] demonstrate that a conditional 1-gram model can be efficiently represented by two-layer (i.e., two attention layers), single-head transformer. In contrast, Sanford al. [29] show that a one-layer, single-head counterpart cannot solve this task unless its hidden dimension is exponentially larger than that of the two-layer model. On the other hand, for higher-order processes, Edelman al. [10] and Nichani et al. [20] argue that low-depth transformers need the number of heads scaling linearly in k in order learn htth -oder Markov processes. However, Rajaraman et al. [26] establish a surprising result that three-layer, single-head transformer can represent the conditional k-ggrm model for any k > 1, and thereby learn 2:th -oder Markov models in-context.
Together, these findings illustrate that while a 1-layer transformer cannot efficiently represent an induction head, 3-layer model with head per layer suffices for all Markov orders k > However, these results leave open a natural and important question:
Can a two-layer, single-head transformer learn :d Markov processes in-context?
In this paper, we approach this question from two perspectives dynamics. In particular, we affirmatively answer this on the representation front and make partial progress on the learning dynamics through the following key contributions:
power and learning
Main Contributions.
1. Improving upon the best known three-layer construction, we prove that a two-layer single-head transformer is sufficient to represent the conditional k-gram model, and thus can learn 4t -odderder Markov chains in-context (Sec. 4). To the best of our knowledge, this is the tightest characterization of transformer depth and Markov order for Markov-ICL.
2. In the course of establishing this result, we uncover an interesting tradeoff between depth and width: the existing three-layer single-head construction (depth is three, width is one) equivalent to a two-layer architecture with two heads in the first layer and one in the second (depth is two, width is two). (Sec. 3)
3. For first-order Markov chains, we prove that gradient descent on our simplified two-layer transformer learns an induction head, and thereby the in-context conditional empirical distribution. (Sec. 5)
Comparison with other Markov-ICL works vis-a-vis transformer architecture. In this paper, as we focus on higher-order processes, we consider the general transformer architecture (Subsec. 2.2)
closer to real-world models with relative positional encodings [30], softmax attention [32], MLPs, and layer normalization [2]. For first-orde processes, past literature has considered simplified variants: Bietti et al. [5] use frozen and fixed random embeddings with no MLPs in the first layer, Edelman et al. [10] consider attention-only models with no MI.ss and embeddings, and Nichani et al. [20] study a variant called disentangled transformer, splitting the residual and attention streams, with the hidden dimension scaling with depth. While our gradient descent result for first-order processes

---

(Sec. 5) is similar to that of Nichani et al. [20], the architecture is fundamentally different. In Tab. 1, we compare the parameter counts of our construction against alternative constructions. Our results show that, as both the order and sequence length increase, our architecture remains the most compact terms of parameter efficiency.
Previous works
Markov order
Architecture
# Parameters
Bietti et al. [5]
2-layer, 1-head
9d2 d$ + Td
Edelman et al. [10]
2-layer, 1-head
215 + 37'S
Nichani et al. [20]
&
2-layer, k-head
(k + (k + (S + T)² T(S TT)
Rajaraman al. [26]
3-layer, 1-head
15(68 + 3)² + (6S 33T 25 10)
Ours
2-layer, 1-head
9(65 + 3)2 + (65 + 32 25 + 9)
Table I: Comparison with prior works transformers with Markov-ICI. Here k >= 1 is the Markov order, T the sequence length, and S the state-space (vocabulary) size. Note that Bietti et al. [5] do not explicitly state what d is and assume it be large enough. To the best of our knowledge, our work obtains the tightest known characterization of transformer depth and Markov order for higher-order processes. The bit precision for our representation results is the same as that of Rajaraman al. [26]: it suffices to have (T) + k) bits per parameter, with (1/"T) additive approximation error. detailed breakdown of the parameter count computations can be found in App. I. Additional experiments analyzing the effect of varying bit precisions are presented in Subsec. H.4.
Notation. Scalars are denoted by lowercase italic letters (e.g., I, y). vectors by lowercase bold letters (c.g., y),,, and matrices by uppercase bold letters (c.g.. A, B). The matrices 0mxn and Imxa represent the all-zeros and all-ones matrices mx respectively. For norms and inner products we use standard notation, i.e., norms written 11. with ||x|| denoting the Euclidean norm. The inner product in Euclidean space between two vectors x and y is denoted by (x,y). The indicator function is denoted either by I(-) or 1()) depending on the context. The vocabulary S= (1, ....,....... S) of cardinality [8] S denotes the finite state-space of the Markov Chains studied in this paper. A sequence Tp:r is defined as Zn:T == Tg. 21.... Dr. where T E N. For any x E S, its one-hot embedding is denoted by e TRSS ((-) denotes the big O notation.
1.1 Related Work
In recent years, transformers have been widely studied from both theoretical and mechanistic perspectives [35, 14, 19]. Foundational results [37, 25, 33] have established their universality and Turing-completenee with further work examining their ability to model formal languages [4. 16] and implement algorithmic behaviors such as induction heads [22]. Concurrently, recent studies have explored how transformers perform ICI., including their learning dynamics and phase transitions [31, 34, 8, 27, 3, I, 11, 11,1,,,, 15]. On this front, our work is most closely related to recent works studying Markov-ICI, where the sequential inputs are modeled as Markov processes. In Bietti et al. [5] and Edelman et al. [10], the authors discover stage-wise learning procedure for first-order sources, and show that a 2-layer, 1-head transformer can represent a first-order induction head. Makkuva et al. [17, 18] study the optimization landscape and learning dynamics of 1-layer, 1-head model for first-order global Markov chain, unveiling the significance of weight-tying and Markov switching probabilities. Nichani al. [20] studies the learning dynamies of 2-layer, 1-head disentangled transformer, illustrating it learns to implement a first-order induction head. For higher-order sources, Nichani et al. [20] constructs 2-layer, k-head architecture, whereas Rajaraman et al. [26] shows that a l-yee,, 1-head model suffices to represent a conditional k-gram model. In contrast, we show that 2-layer, 1-head transformer suffices to represent arbitrary-order Markov models in-context.
2 Background
In this section, we provide the requisite preliminaries on Markov processes, the conditional k-gram model, and the Transformer.

---

2.1
Markov Processes and the Conditional k-gram Model
Input:
D
A
Q
D
A
D
Unifarm Distribution
Attention:
Input:
Match last k symbols
3
7
(b) k -order induction head (Def. 4) The attention (a) The conditional k-gram model (Def. 2). It first pattern concentrates positions in the sequence where (1) identifies the positions the sequence where the preceding & tokens match the current context (bluc), These are the positions the model attends to most and (2) returns the empirical distribution over the strongly, as they reflect the same context as the current symbols (green) that follow these matched positions.
preceding tokens exactly match the final tokens.
prediction target.
A hmhh -oder Markov process is a time-homen humgenoogssssuu stochastic process over finite state space, where the probability of the next symbol depends only on the most recent & symbols in the sequence [21]. More formally, it is defined as follows:
Definition 1. (hth -oder Markov Chains) A stochastic process žo:T == Io, T1,-.. over finite state space S is said to be -oder Markov Chain with transition kernel :(-) if, for all & & & T, the following condition holds:
Fa In. mo) x( = En, In In k+)) V.
We now define the empirical equivalent of the kernel æ, the conditional k- k--ramm model:
Definition 2. (Conditional k-gram model) Given sequence Fo:r == To, 21," in ST-1 the conditional k-gram model fxx(s Ar, probability of observing a symbol E S matching the last k tokens, i.e.
To) computes the empirical
I( 34-1 IT, Fi-2 TT-1. Fi-k k FT -T+1+1)
(x Zo. Z1, ,r)
I( NT, 2i-2 -1....... Ii xT-K+1)
assuming the denominator is non-zero. Thus conditional k-gram be interpreted simple count-based estimator the transition kernel π,, using the in-context count estimate for the predictive probability distribution.
Data generation (Random Markov sequences). To empirically generate the input data, i.e. random Markov sequences of order k. we first sample each row of the transition matrix я independently from a Dirichlet prior with parameter vector 1, which is equivalent to an uniform distribution on the S-dimensional probability simplex. The initial & tokens of each sequence sampled uniformly from S and remaining tokens are generated according to the sampled kernel æ. To sample an input batch, first sample multiple independent transition kernels, generate sequences from each, and then aggregate these sequences to form a batch for training or evaluation.
2.2 Transformer Model
While there exist several attention- transformer architectures in literature, we adopt the formulation described in Shaw et al. [30] and Dai al. [9], which employs relative positional encodings. Each layer in our transformer comprises softmax-based self-attention, optionally followed by a multilayer perceptron (MLP). Following the construction outlined in Rajaraman al. [26], we define L-layer transformer with multi-headed attention follows:

---

Definition 3. (Multi-head Attention Transformer)
Require: Input T0.7, number of layers L, number of attention heads per layer H
Ensure: Output distribution Po(- (o:T)
x Emb( for n 0,l,...,
for = 1 to L do
for n to T do
for h = 1 to He do
((.bead:h) X
Ei= att,
(l.hoadsh))
(w/fthoad:
(2) X + Pn-i
(f.bead: h,
V))
end for x (((+)) = Xn
+ - Xn (ff.haad:h)
> Residual connection
((+1) MLP(x/) , xx(+1))
MLP with layer normalization and skip connections
end for end for
logitsy = W ox xxx (L+1) + bo
Po(: TO:T) :f::==((ppp (loggtss))
Where the attention weights are defined as:
(Chead:hh,K) (Chhea::h)) = sof tmax, (< (ww(haeesee (Ehhed:hh) + P - n,i
(6.bead:/)) (8)) x
att
W
Q
Here denotes the layer index and h the index of the attention head within a layer. The matrices w/ wff.heeash) w (h.head:h) w-ff.bad:h)) E IRdxd are the query, key, and value projection matrices, (f.bead:h.V) E IRd represent the relative positional respectively The vectors Pn-; P (C,head:h.K) and Pa-i
Q
K
encodings that modulate attention based on token distance. Finally, the output the transformer is projected using W, ((xx(++) + b, IR$, mapping the representation from the embedding dimension the output vocabulary size S. Although softmax is typically used for output normalization, in setup we define ʃ(.) ReLU(-). Note that the attention mechanism described is causal self-attention.
3 Warm Up: Construction with Two Heads in Layer One, One in Layer Two
represent ttth order Markov processes with a 2-layer, 2-head transformer. Towards the same, we recall that the 1.h in-context counting estimator, i.e. conditional k-gram model, defined in Subsee. 2.1 is closely related the classical Laplacian smoothing, which in turn is the optimal Bayes estimator for the next-token predictive distribution [12, 28]. In view of this fact, in order estimate the optimal conditional distribution of each token in-context, it is natural to ask: how does a transformer represent this conditional k-gram model? To this end, Rajaraman et [26] introduces the notion of k-th order induction head. This mechanism enables the attention layer of a transformer to assign at each time step the highest attention weight past tokens whose length-k context matches that of the current token. For the sake of completeness, we recall the formal definition.
As up, in this section, prove our first main result that
Definition th order induction head) An attention layer with a single head is said to implement hthh order induction head if, for any input sequence (æo, ,1,..,,,,,77 (T) ST+1 and for any fixed index i & T, the attention score atty,; at position is maximized if and only if its preceding & tokens exactly match the final tokens of the sequence. That is, attr, is maximized and only æi-j = IT-j+1 for all i (1,2,...,k),
The salience of this higher-order induction head is immediately reflected by the fact that such mechanism is indeed what precisely needed to implement the conditional k-gram estimator from Def. 2. Capitalizing on this, Rajaraman et al. [26] proves the best known result for higher-order Markov processes, through 3-layer, 1-head transformer architecture, which we recall below to motivate main result:
5

---

Theorem 1. (Theorem 4 in Rajaraman et al. [26]): The conditional k-gram model can be represented using a transformer with three layers, each containing a single attention head. The layers separated by MLP blocks and include relative positional encodings well as layer normalization. The embedding dimension scales as (SS).
Proof sketch. We provide a brief outline of the proof of Thm. 1. For each position n in the sequence So:r, the first attention layer effectively learns a hard attention map concentrated on the indices {m, n ".....,7. 72 k + 1]. Specifically, within a subspace of its embedding space, this layer (2-00 with attention weights abbmn,, 3*/ ())00))
computes u, - 3/ -
for i E {n, n I,.... past-k context starting at position n. Similarly, the second attention layer focuses on the positions capturing the context shifted by one position. After these two layers, the embedding at position n contains both ///Ivull2 and the normalized vector u,,/ u,lll2, assuming appropriate choices of embedding dimensions, query/key/value matrices, positional encodings, and MLP layers with layer normalization. In particular, the architecture can configured to compute the f2-norm via layer norm, as detailed in Subsec. A.1. The third attention layer then functions as a k-th order induction head: at position n, the attention score made proportional to the cosine similarity (ur, vn)// | ||//| v,ll2))) This quantity equals I when uy Yn. and is strictly less than I otherwise. As the softmax temperature tends to infinity, or equivalently, as the attention logits are scaled by a constant tending to infinity, the model approaches the behavior of conditional k-gram estimator. Hence, the third layer effectively implements a k-th order induction head.
1), and zero otherwise. Intuitively, this u, allows to capture the
(n 1,n 2.... k), and computes Vn = (Et - e/..)/ (3"-
We are now ready to prove our first main result. Specifically, we show that the aforementioned construction with 3 layers and a single head can be adapted to a two-layer counterpart with two heads in the first layer and one head the second.
Theorem 2. (Two heads in the first layer, one in the second): The conditional k-gram model can be represented using transformer with two layers, the first containing attention heads and the second containing a single attention head. The layers are separated by MLP blocks and include relative positional encodings as well as layer normalization. The embedding dimension is 65 + 3.
Remark 1. (Depth-width tradcoff) Thm. thus reveals an interesting tradeoff between depth (number of attention layers) and width (maximum number of attention heads per layer). While our construction here has both depth and width two, its counterpart in Thm. I instead has depth three and width one. Thus our new architecture trades off depth with an additional attention head, thereby increasing the width.
Proof sketch. We provide a brief sketch of the proof for Thm. 2, using the same notation as that of the Thm. 1, and focusing on the key differences between the two architectures.
Layer One. For each position n in the sequence su.T, the first attention head in the first layer learns
focused attention pattern over the positions (n,, n 1.......... n n k 1). a subspace of the embedding space, this head computes the representation u, 2: 3 - .e.... --=--- 3:) On the other hand, the second attention head in this layer focuses on the preceding context, attending to positions {n - L, 2...., n- - k). and computes the representation Vn, = (> 3¹. - /
/>
3))
After this layer, the embedding at position contains both the vectors un/ |ulll and Vn, f2/11vn.1.. achieved by choosing an appropriate choice of embedding dimensions, query, key, and value matrices, positional encodings, and MLPs with layer normalization.
Layer Two. The second attention layer now plays the role of an induction head, analogous to the third layer in Theorem At position n, this head computes an attention score that proportional to the cosine similarity (ur, va>/[lur v,ll2. This score attains the maximum value of one when

---

ur Yn, and strictly less otherwise. As the temperature parameter in the softmax becomes large, or equivalently as the inner product is multiplied by large constant, the attention mechanism approaches the behavior of a conditional k-gram estimator. Therefore, this two-layer transformer, with two heads in the first layer and a single head the second, sufficient to represent the conditional k-gram model. We refer to App. A for the full proof.
4 Main Result: Two-Layer Single-Head Construction
Learnt Attention Map
Learnt Attention Map
Optimal Attention Map
Layer
Layer 2
Figure 2: Attention maps learnt by two-layer, single-head transformer trained sequences generated by random Markov chains of order 3 (Subsec. 2.1). (i) the first layer, the attention map shows clear pattern: attention weights increase monotonically along the first three lower-diagonals and drop to zero beyond that. This suggests that the relative positional bias is maximized the diagonal with an offset of kk -3, i.e., the third diagonal below the main diagonal, which is consistent with the construction in Sec. 4, (ii) in the second layer, the attention map closely resembles the ideal attention pattern required to approximate the conditional k-gram estimator. We note that all experiments were conducted using standard initialization schemes. Additional experiments with different orders of markov chains and experimental details are provided App. H. In Subsec. H.2, we also experimentally demonstrate that single-layer transformers fail to solve the induction head task with the same order of parameters. Finally, Subsec. H.3, we test the robustness of the two-layer, single-head model to noise in the input sequences.
While Thm. 2 above demonstrates that the conditional k-gram can be represented by two-layer transformer, however relies on a two-head construction, thus still leaving open our motivating question: can 2-layer, 1-head transformer represent the conditional k-gram? In this section, precisely address this gap and realize 2-layer, 1-head model to represent lo"-order Markov processes for any k > 1. The key idea behind our approach is to leverage the MIP--and specifically, the non-lincarities such as ReLU and LayerNorm-4e to isolate symbols at particular positions that are useful for estimating induction heads. This contrasts with prior constructions, which primarily focused on the attention mechanism while underutilizing the role of the MLP and its non-linear components. As result, our construction highlights how these non-linear MLP elements can play a critical role in enabling in-context learning, demonstrating their significance. We now present our main result.
Theorem 3. (Two-layer, single-head construction): The conditional k-gram model can be represented using a transformer with two layers, both containing a single attention head. The layers are separated by MLP blocks and include relative positional encodings well as layer normalization. The embedding dimension is 65 + 3.
Remark 2. (Parameter count and bit precision) We note that the above construction utilizes 966S + 3)2 + (65 + (27 + 25 + 9) parameters. A slightly more efficient construction in terms of parameter count with 8(65 + 3)² + (65 + (2T 25 + 9) parameters can be found in App. B. For both these results, suffices to have (lo + k) bits per parameter, with (L/TT) additive approximation error, where 7 is the sequence length.
Proof Sketch. We now provide an outline of the proof of Thm. 3, using notation consistent with the constructions in Thm. and Thm. 2.
7

---

Layer One. For each position in the sequence Su:T, the attention head in the first layer attends to the preceding context, focusing on positions (n nn-1,nn-2.... n- and computes the representation
Vn 1 3 e....) ( with attention weights atin,,i 3 (Σj-133 3 for
{n - ,,.... k), and zero otherwise. This matches the optimal attention pattern for Layer-1 illustrated in Fig.
MLP. In contrast previous constructions, the MLP situated between the attention layers plays a central role in our proof and is structured to reconstruct the complementary representation u, = ) which corresponds to a hard attention pattern over the positions
(2x-0 36 eS - )/ (2 {n, n- 1, ..... n2 k 13. Recall that this un was constructed Thm. 2 thanks second attention head, absent in our architecture.
Our main idea tackle this issue is isolate the one-hot embedding c from vn computed in the first layer. This is achieved by embedding the entire one-hot vocabulary within the weight matrix of the first MLP layer. When combined with a suitable bias and followed by a ReLU activation and layer normalization, this configuration filters out all but the component aligned with Zn-k. thereby producing e = At this stage, the embedding also retains the representation Vn, computed by the first attention layer. Subsequent layers, along with appropriately configured skip connections, implement a linear transformation that combines ex - Vn, and the token embedding ei. the latter of which is introduced via a skip connection. Since v, and #n_ occupy orthogonal subspaces of the embedding, and all required components are explicitly available, this construction yields the intermediate representation u, (c Ex=13) 3+)= (3e (3e5... (E--033 3r)
3* eg C / (E 3 ))
which simplifies exactly to the desired form u, - 3 appropriately chosen weights, layer normalization, and skip connections, the resulting embedding at position encodes both un / and vn/ vxll2-
/E/-0 13)))) 36) With
Layer Two. As in the second layer of Thm. 2, the attention mechanism in the second layer is configured to act as a hdh induction head by computing an attention score at position n that is proportional to the cosine similarity (uu,, vx>/||rr ////////l// a 2). This score reaches its maximum value of one when uy = Vn, and is strictly less than one otherwise. As the softmax temperature tends to infinity, equivalently the attention logits scaled by constant that tends to infinity, the resulting distribution converges to that of a conditional I- k-gram estimator. App. B details the complete proof.
Importance of non-linearities. This construction demonstrates that non-linear components such as ReI.U activations and layer normalization are not merely auxiliary, but play a critical role in enabling the transformer architecture to express higher-order inductive structures that essential for in-context learning.
Note. Due to the problem's 's non-con non-convexity, multiple constructions may enable two-layer, single-heat transformers to implement a 1tth -order induction head; our approach offers one such construction motivated by two key insights. First, empirical observation: for datasets generated from 4h -odeer Markov chains, we observe that in trained transformer model, the first-layer attention weights increase monotonically along the first 1 lower diagonals before dropping sharply to near-zero (see Fig. 2 and Subsec. H. Second, prior theoretical construction: as demonstrated in Rajaraman et al. [26], the attention pattern can be designed to be proportional to a dyadic sum representation. Taken together, they offer both empirical and theoretical grounding for our construction of -odder induction heads.
5 Gradient Descent Analysis
Thm. 3 establishes a representation result showing that a two-layer transformer with single attention heads can implement the conditional k-gram model. However, this does not guarantee that current optimization algorithms such as gradient descent (or its stochastic variants) can recover such a solution in practice. To this end, in this section we focus first-order Markov chains and show
8

---

that gradient descent on reduced variant of our two-layer transformer indeed learns an induction head, and thereby the in-context conditional 1-gram. On this note, we would like to emphasize that first-order Markov analysis a critical first-step and a building block for challenging higher-order analysis [20].
Building towards our result, we quickly recall a few recent works [10, 20] that also analyze the gradient dynamics on first-order Markov data with simplified transformer architectures. Specifically, Edelman et al. [10] study an attention-only linear transformer without MLPs and softmax operations, while Nichani et [20] analyze disentangled transformer, with residual and attention streams split. On the other hand, our transformer architecture in Def. 3 consists of relative positional encodings, softmax attention, MLP with ReLU activations, and layer normalization. While non- non-linearities such as ReLU and layer normalization play a crucial role in our representation result, as illustrated in Sec. 4, at the same time, they make the gradient analysis challenging and intractable even for simple first-order Markov setup.
To tackle this, we study a reduced model wherein we only treat the salient components of our two-layer architecture in Thm. 3, such as positional encodings and attention layer, as parameters, and treating the rest as approximately optimal. This enables for a a faithful reproduction of the ICL phenomenon exhibited by the parent model, whilst being tractable. This is akin to the reparameterization strategies of Makkuva et al. [18] and Nichani et al. [20] for gradient-flowgrradee how//rrdienttdeseeeee Alternatively, this can also be viewed as a form of good parameter initialization, similar to that of [10, 20]. We now start with our technical assumptions on the data distribution and the transformer architecture. We use notation from Subsec. 2.2, dropping the Seaddddentifying superseript.
analysis.
Assumptions.
(i) Data distribution assumptions: Following [20], we assume the prior over transition kernels enforces non -onneeenerr,y, positive transitions, and constant mean (see App. C).
(ii) The positional embeddings used for the keys in the first attention layer are scalar-valued, i.e.
(1,K),
Pn
Pa. where pn is a scalar.
(iii) ((1) Q.K,V and p(1.V) are chosen that the attention weight for position is given
n. exp(p:)/ exp((p))) E {0, 1), and the corresponding value output is V, Ei_ att 1/ "
by att(1) (1)
(iv) We assume the MLP is ppiiall-ie... it outputs the correct un even from suboptimal Yas this holds for first-order Markov chains via skip connections.
(2)
Q.K are chosen such that the attention in this layer is defined as at!
exp (a2 uT)) (o exp(02 (V;, (v)), where 02 E lK is the attention scalar.
(2) Ex, ER. (vi) w (2) V is chosen so that the final logit logit,- L.O att,
Loss function. We minimize the cross-entropy loss for next-token prediction [20]:
C(0)- -E.TPP. ST
Σ (aT+1 xo.r) log (logitz (s) + e)
(I)
where E > 0 is a small additive constant, Px (-) denotes the prior over Markov transition kernels, and (sg is the logit corresponding to the symbol s.
Training algorithm. We denote the of trainable parameters as 0 (p. az), where the positional scalars p [Po, .i.... | and ag is the attention scalar. Similar to [20], we adopt a two-stage training procedure. (i) In the first stage, we optimize the positional scalars p using gradient descent for Tj steps with learning rate 7/, (ii) in the second stage, we freeze and train only 42 for an additional T2 steps using separate learning rate 7/2- key distinction in our setup is how we treat layer normalization during training. While in the first stage we omit all layer norms in the MLP, during the second step we reintroduce them. This stage-wise enables a tractable theoretical analysis. We refer to App. C for further details. We now present main result:

---

Theorem 4. (Convergence of the training Algorithm): Suppose that Assumptions (i)- (vi) hold and that the Markov sequence length I' satisfies T+ 1 > poly(7 ,,S) for some constant Y O. Then there exist 0, learning rates 7, 72. and step counts T1, T2, such that the output of the above two-stage training algorithm, 0 (1p: (2), satisfies
c(0) L* log( (T' 1) (7 + 1)ey
for constant c independent of (), and the optimal loss -B (1/S) ( s) log (s' 8)
C
Thus Thm. 4 illustrates that the transformer model trained via the two-stage algorithm achieves near-optimal loss asymptotically. Furthermore, we can also show that & approximates the conditional 1-gram model, i.e. first-order induction head, with vanishing error as sequence length grows (see App. C). For completeness, perturbative analysis of the architecture consisting of two attention heads in the first layer and in the second (see Sec. 3) for any-order Markov chains can be found in App. D.
Conclusion
In this paper, improving over prior three-layer constructions, we show that a two-layer, single-head transformer can represent any k-th order Markov process. Thus our result gives the tightest known depth characterization for in-context learning conditional k-grams. Additionally, we provide gradient descent analysis of how such representations emerge during training first-order Markov data. While these results deepen our thcoretical understanding of ICL and show that compact transformer architectures particular, fully characterizing the learning dynamics for higher-order processes is an important avenue of future research. Further, we view contributions a first step toward understanding the fundamental limits of ICL in compact transformer architectures, offering not only tractable setting for theoretical analysis, but also potential pathway toward more efficient model designs.
be both expressive and learnable, several open questions remain. In
Acknowledgements
This work was supported in part by the Swiss National Science Foundation under Grant No. 200364. We are grateful Eshaan Nichani for insightful discussions and valuable input throughout the preparation of this paper.
10

---

References
[1] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2023.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL https://arxiv org/ aaaaab///// /tbb/107..o 0w1116670065655
[3] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.
[4] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 7096-7116, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v 10055///1122220mmmmmmma......35 org/2020. http//220.emmmpmmmaa... 576/.
URL https://
s.llattoooogy.org222222
[5] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer. memory viewpoint. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:t//opeeeeeerie http://ewee......ooooooadaa/a//.
[6] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
[7] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.10555 2024.
[8] Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts to function lipschitzness. Advances Neural Information Processing Systems, 37:92638 92696, 2024.
[9] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context, 2019. URL https:
//arxiv. wwwp/ogg///abbeeeeeeeee abs/1901. http:/9911...0.00.1
[10] Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. Advances in Neural Information Processing Systems, 37:64273-6433 2024.
[11] Jesse Hoogland, George Wang, Matthew Farrugia- Roberts, Liam Carroll, Susan Wei, and Daniel Murfet. The developmental landscape of in-context learning. arXiv preprint arXiv:2402.0224 424.0 2024.
[12] Pierre Simon Laplace. Essai philosophique sur les probabilités. Courcier, Paris, France, 1814. Reprinted by Cambridge University Press, 2009. In the reprint, the estimator appears page 23.
[13] David Levin and Yuval Peres. Markov chains and mixing times, volume 107. American Mathematical Soc., 2017.
[14] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pages 19689-19729, PMLR, 2023.
[15] Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning, 2024.
[16] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shorteuts to automata, 2023,

---

[17] 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.
[18] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Hyeji Kim, Michael Gastpar, and Chanakya Ekbote. Local to global: Learning dynamics and effect of initialization for transformers. Advances in Neural Information Processing Systems, 37:86243-866 2024.
[19] Timothy Nguyen. Understanding transformers via n-gram statistics. Advances in neural information processing systems, 11998099-98088 2024.
[20] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning, 2024. URL attps://openrrevess wwtp:/iwwweieeeeeeeeeeeeeeeeeeee wwt.forrm htrrmoorrmmm id= jimmiiiiim12H22
[21] James R Norris. Markov chains. Number 2. Cambridge university press, 1998.
[22] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 2022.
[23] OpenAL. Gpt-4 technical report, 2023. URL https: openao.pena..... wwt.coo.//appcc com/ppapessgpt-44.. pdf.
[24] Matteo Pagliardini. GPT-2 modular codebase implementation. https://github: wtt.:/cccccccpeeeeeeee /pfml/ 1lm-baselin ines. Accessed: Jan. 2025.
[25] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing-complete. Journal of Machine Learning Research, 22(75):1-355 2021.
[26] Nived Rajaraman, Marco Bondaschi, Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar. Transformers on markov data: Constant depth suffices. Advances Neural Information Processing Systems, 37:137521-137556, 2024.
[27] Ruifeng Ren and Yong Liu. Towards understanding how transformers learn in-context through representation learning lens. Advances in Neural Information Processing Systems, 37:892-933. 2024.
[28] J. Rissanen. Universal coding, information, prediction, and estimation. IEEE Transactions on Information Theory, 30(4):629-636 1984. doi: 10. 1900///T.1994.1055655..
[29] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail solve the induction heads task. arXiv preprint arXiv:2408. 1110008..4333 14332, 2024.
[30] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint 0ttxXvv/800302/5555 2018
[31] Aaditya K Singh, Ted Moskovitz, Felix Hill, Stephanie CY Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. arXiv preprint arXiv:2404. 07 129, 2024.
[32] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[33] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. In Advances in Neural Information Processing Systems, volume 35, pages 12071-12083, 2022.
[34] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24877, 2022.
[35] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference Machine Learning, pages 11080-11090. PMLR, 2021.
12

---

[36] An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint arXiv:24/2.55133 15115, 2024.
[37] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sqqueee functions? International Conference Learning Representations, 2020.
13

---

Appendix
Contents
1 Introduction
1.1 Related Work
2 Background
2.1 Markov Processes and the Conditional k-gram Model
2.2
Transformer Model
3 Warm Up: Construction with Two Heads in Layer One, One in Layer Two
4 Main Result: Two-Layer Single-Head Construction
5 Gradient Descent Analysis
6 Conclusion
10
A Two Heads in Layer One, One in Layer Two Construction A.1 Modification the Layer Norm
16
16
A.2 Construction
16
A.2.1 Two Layer MLP followed by Second Attention Layer
18
A.2.2
Layer Norms in the Second Attention Block
20
Two-Layer Single-Head Construction
23
B.1 Three Layer MLP followed by the Second Attention Layer
.ay.............................
24
B.2 Two Layer MLP followed by the Second Attention Layer
................................
28
Two-Layer Single-Head Gradient Analysis
33
C.1 Model Simplification
35
................................
C.2 Stage Analysis
................................
37
C,3 Stage Analysis
41
D Two Heads in Layer One, One in Layer Two Perturbative Analysis
42
D.I Construction
43
D.2 Convergence Analysis
45
D.2.1 Warmup: Second Order Markov Chain
48
D.2.2
Any-Order Markov Chain
51
E Training Dynamics: Stage One Lemmas
53
F Higher Order Markov Chain Lemmas
60
G Additional Lemmas
64
H Experiments and Experimental Details
67
H.I Additional Experiments
67
H.1.1
First-Order Markov Chains
67
H.1.2
Second-Order Markov Chains
68
H.1.3 Third-Order Markov Chains
69
H.2
Experiments on One Layer Transformers
70
H.3
Experiments Noisy Sequences..
71
H.3.1 Random Token Substitution
71
H.3.2
Perturbation of Transition Dynamics
71
H.3.3
Discussion
72
H.4
Experiments on Bit Precision Experimental Details
72
H.5
72
14

---

I Parameter Breakdown and Comparison
74
15

---

A Two Heads in Layer One, One in Layer Two Construction
A.1 Modification of the Layer Norm
Following Rajaraman et [26] (Appendix C.I in Rajaraman et [26]), we use modified version of layer normalization extract the Eg-normalized direction of the embedding Xn. Specifically, augment the embedding by concatenating it with its negation to form [x.: xX,], which ensures that the mean of the features is zero. Applying standard layer normalization to this vector yields vd x.. |xall in the first components, since the standard deviation of the augmented vector is Xx,, lolllvd d. To recover the unit-normalized embedding x, //xl/ simply divide the first components of the output by vd, and extract the first d components. Throughout this work, we assume this normalization and scaling step is handled implicitly, and we interpret layer normalization as returning the normalized embedding direction. In order to avoid explicitly repeating this transformation at each layer, we redefine the layer norm to implicitly return v/ v/||v|22 for any input vector v. Hence, v/Iv/2 = LV()). where, (.)()) denotes the modified layer norm operator. Note that this modified definition used throughout this paper.
A.2 Construction
(Two Layer, Two Heads in the First Layer). We claim that there exists a two-layer attention- based Transformer model fe. consisting of two attention heads in the first layer and single attention head in the second layer, which can be configured in one of two equivalent ways: (I) with two MLP layers inserted between the first and second attention blocks, or (2) with no intermediate MLP layers, but with layer normalization applied the query and key projections in the second attention block, such that:
Theorem 5.
fo(BU.T) < 81.7) (SI.T 8') n, Pa.
Note that Px denotes an arbitrary distribution over trajectories, where each trajectory is generated using a Markovian transition kernel The kernel corresponds -oder Markov chain.
Importantly, in both configurations, the input embeddings and the structure of the first layer (Layer 1) remain identical. We define the input embeddings as follows:
Emb( nx)
[0 x3
01x5s]"
TR65+
In the first layer, for both the heads use the following relative positional embeddings
k [1 01x(1+65)] ]"
if 0, if i E {1, 1 1},
(head: 1).K
(head:2). :2),
(i log(3)) + x)) [o
0 01x(1+65)]
P
P.
( log() k:) [o 01x(65)
if i k, if
and the same, value embeddings.
().
(head:2).V
3
0)"
for i $ k i > k.
p
Layer 1, Head 1. Consider the key and query matrices,
W
(kead:1)
lix2
01x1
01x6s (22465)x(+455)
K
0 (2465)
(2465)x1
(haad:1)
01x3
1ixs
0xxss
Q
0
(24055)x3
(2+G5)x5
(2+65)x55
16

---

On computing (for i to = 1)
w (head::))
Emb( (
W (head:
(haad:1)::K [i (i. log(3) + k) ((465)x1]
K
W (head:1 Emb((In) = [11x1 0 01x(2+07] ]
Then, observe that,
W (1)aa:11 (Emb(In-s) ( W [head:: Emb Emb((2n) (2n)) (e log(3) + k) 11(0 < ii min k 13)
K
Letting - c0, this results in the attention pattern,
att (head:1) n,n-
3'1((0 & is min 13) 3'1(0 & i min(n, 11) 3'1(0 & & min(n, A 13)
-0
mmnn 3
-0-
mminin.8)
3
Catta
Choose the value matrix as,
03x3
03x
03x4S
(bead:1)
035x3
0sxs
038x
0sx3
TsxS
0sx4s
028x3
Osx
028x48
Hence, the output of the attention head is:
X [bead: :1
[
| 01x2
0ixs Un
01x2s
0 01x2s 17.
min(n.k-1)
where, Un
at a Can
e
Layer 1, Head 2. Consider the key and query matrices,
W (head:2) K
01x1 0
1ix2
01x6s
(24655)
(2+05))x
0 (2+65)
W (haad:2)
01x3
1ixs
0 0xx58
0(2465)x3
((4465)
(2+65)x55
On computing (for i = 1 =
W Chead:2 2)
Emb( Xn =
K
(head:2) (1).K = [(i log(3) + k)
0 T
p
W (head:2) Emb(xn)) = [l1x1 0 01x(2+657)
Then, observe that,
w ()) K
'Emb(rr (In-i) p
(bead:2).K) w whheod: (2ead:2) Emb (Im)) (i- log(3) + x) 11( <is min(n,k)) k3)
Emb(xx))
17

---

Letting - co, this results in the attention pattern,
311(1 & & min(n,k)) n,n-i Emin(n,k) 3
I111 & is min(n,k)) 3'-1T(1 & i & min(n, kj)
att, (1)
Ex=o
min 3
Cataa
Choose the value matrix as,
13x3
0
03x45
0asx3
Osxs
0a5x4S
sx3
Isxs
0sx4S
028x3
Osxs
025x45
The output of the second attention head is,
X
(hand:
=
Zn
01x2
0ixs
01xas
Vx
0 x22
1
min(n,k}
where, Vn
Σ
attn,n- es
xm-
min(k.n-1)
Zn
Σ
atl,,n-i3
is straightforward to verify that Zn 3k+1/5 when k 1, and that Zn < 3*/5 otherwise. This observation will be useful later, as the value of Zn serves as a signal for whether 3 k + 1 or n & In particular, this distinction enables the next layer to bypass attention computations for indices i & k, where the condition I _
fa- = Ti-k is not well defined.
Therefore, with skip connections:
(1))
(besd::)
x
(head:2)
Zm
1x2
Un
01xs
Vn
01x2s J"
A.2.1 Two Layer MLP followed by Second Attention Layer
MLP.
The first layer for the MLP
0 (3+3))
(3+5))
(3+ (2s))
(1
(3+388)
OsxS
0 0sx25
mlp
(x(33))))
Isxs
0x22
0 (3+33)
03sx5
0 3Sxx28
0 (3+65)x
mlp
Incorporating skip connections and applying non-linearities, we obtain:
((1)
LN (ReLU
(2) (w!()>()))
baip)
nmmpp
(I)
Z,
01x2
ex
un
|
01xs
01xs
The second layer for the MLPis:
18

---

0 (3+33)
(3+
(3+8)x(22))
0 (x+339))
8xx
w: (2)
mlp
(x(3438)
0sxS
Sx28
(3+-38)
Isxs
xx
0 2 (3++35)
028x8
2Sx28
(2)
mlp
(3+65)x1
Including skip connections and non-linearities, get:
X
n.mmpp
(2)
xn,mlp
(1)
+ LN (ReLU
(
(22)) MipX
+ Dmip
m.p
(2
Zn
1x2
V/
1xs
Layer 2.
this layer, all relative position encodings are set to O, and we use the following query and key matrices for the attention mechanism in the second layer:
w Q
VK
(2+33)))
Is
3
VA
(2+444)))
Isxs
With these choices place, we obtain:
(2)2222)
(Wl2) (2) K ) (mmip)) ZiZn +
R. Ov: Un)
K 7i7n + 2 k
V;
Un
Vell2 |ul122
|v
u,
Applying the softmax function, we obtain:
exp (kZzZn 2
1vil2
Iu,, la
att n,
Σ jー exp (xZ, 2n + 2x
un
15,12
We can see that for all k+1 < that of position n, i.c., {I;-1-j æn=j for all = O,..., k}. Moreover, for any & + i n, Vi o u. then the normalized vectors are separated by non-negligible margin:
u, and only the local context around position i matches
V: Ivell2
1
Un,
2
2
36
Note that while this gap is small, it strictly non-zero. Recall that Z. 2 for i > and
Zi < gk
5 otherwise. As result, the attention mechanism favors positions i such that U, and the limit as - cc, this results in the following attention pattern:
att (2) n,
Unif(Ix),
19

---

where In denotes the set of all indices E {A + 1,..., ,,] whose local context matches that of position n, formally defined as:
In = (ki -t In-j for all 0, o,....k)
Moreover, (T,) denotes the uniform distribution over the indices in Zn; that is, each position i In receives equal attention weight, and all other positions receive zero. We now choose the value projection matrix in the second attention layer as:
W
0sx3
I
0
With this choice, the output of the second attention layer becomes;
(2)
(2)
22)
Xn,mlp
att
n2
Exrq
S
Xn.
In
For the output linear layer, we choose:
W
(Sx((S+3)
[sxs];
ba 0.
This yields the output logits:
I(V k, Ki-j In-j+1) Fr 1 k. zØ-j In-s+1)
logit,
c
In
sEIn
Finally, by adjusting the notation and setting obtain:
T to reflect prediction the final time step,
T(VO In-i - JT
(2T I In-i IT+1+1)
which precisely computes the conditional k-gram probability.
A.2.2 Layer Norms in the Second Attention Block
(1) from the first layer as defined.
We start with
x(1)
(1)
(bead:1)
(head:2)
Zm
1x2
Un
01x
0xxes J"
Layer 2.
In this layer, all the relative position encodings are set as and instead,
W (2)
0
0 Dxx (2+35)
Isx s
0
0
w(2)
0
0
(24-5)
Isxs
0
20

---

With these choices, and adding a layer normalization after the output of the key linear layer and the query linear layer
(NN(W
2x ( (2) L k.mip)
2(v., Un
Vill
|nl12
V
Un
2
Va
Un
1
Taking the softmax with a temperature parameter K,
exp (2x K
Tvcl
att
=j= exp (2k Ivslla
Note that although this gap is small, it remains non-zero. In particular, as 00, the resulting attention pattern approaches
(t = Uni Ix),
where In denotes the set of indices i (k....,n} whose local context matches that of position Formally,
In = f (k..... n] 2t-1-j In-j for all .,...,k)
In contrast, La captures earlier positions aligned with the normalized query at position n!
(O,...,k 1) whose (normalized) key vectors
4 = { ie t0..... 13
v; |/s/|2
kn
Iu
On choosing,
0 0sx3
TsxS
0
We obtain,
(2)
(2)
+
atl
n,
S
(3)
S
Inl
+
I
For the output linear layer, we choose:
W (sx(55+3)
Isxs].
b, 0.
which results in,
[[VI Ej k, Iz-j = xn-j+1)
logit,,
Σ
a4 Σ
I[VI & k, xr In-j+1) + Ix c
Inl
[
E
VeTn of
i-k
21

---

I [M b e
Σ VE
E I(VI EJ k, zz-j In-j+1) + It
Now, by substituting = T, and aiming to predict the next symbol given the current context, we obtain:
Σm 1( si k. In-i = xT-1++)
-m I(VI Ei k, In-i IT-++1) Ix
Eie, (E++1,
_ nーk I( & k Tn-i = ri+1+1) + Ix
This corresponds biased conditional k-gram model, where the bias originates from the second term in the expression above. Notably, this bias vamishes in the limit as T - 00. The reason is that the numerator of the second term remains constant with respect to sequence length, since I, is fixed and does not scale with T. In contrast, the denominator, given by Σ T(V1 < i & k, Tn-i 3r--+1), grows with T, assuming the underlying Markov chain is irreducible (i.c., every state is reachable from every other state). Consequently, under this assumption of irreducibility, the influence of the second term diminishes as T becomes large, yielding the approximation for a sufficiently large T';
ΣT Σm I[V0 Eish k Zn-1 = xT-1+1)
logity (#T+1)
77 2m I[V1
< h Zn-i = 2r-1+1)
22

---

B Two-Layer Single-Head Construction
Theorem 6. (Two-layer, single-head construction). We assert that there exists a two-layer, single-head attention Transformer model fo. which can be configured in one of two equivalent ways: (1) with three MLP layers placed between the first and second attention blocks, (2) with two MLP layers between the attention blocks, but with layer normalization applied to the query and key projections the second attention block, such that:
Jo(S):T) क (5 81:T) 4 (81:T 8') n, π P...
Note that Pxx denotes arbitrary distribution over trajectories, where each trajectory is generated using a Markovian transition kernel m. The kernel π corresponds to a -oh-order Markov chain.
important note that in both configurations, the input embeddings and the first layer (Layer 1) remain identical. We choose the input embeddings follows:
Emb(x,,)
-0
6x 0 01x5s] 655+
In the first layer we use the following relative positional embeddings:
if 0,
P m).K
(i 1og(3) +K)- [O 01x(1+65)}'
i 21,
if k
and the value embeddings,
DD.
3 [1 oj
for i & k k.
Layer 1. Consider the key and query matrices,
w(()
01x1
Lix2
01x6s
K
0
((2+65)x1
(2+65)x2
(2+65))((+665)
(
01x3
lixs
01x5s
(2+655)x3
(2+65)xS
(2+655)
On computing (for i i K)
w "eE((((nn)
W p (1).K ((i-log(3) k) x12-
w(1) Emb((r,)
[lIx1 0 01x(2+657] I"
Then, observe that,
w()) (in p (). w Emb((r)) (i. (og(3)() + x)) T(1 & &iS min(n, R))
Letting - 00, this results in the attention pattern,
23

---

3'1(1 & i min(n,
3 E min(n,k)) 3'-1I(1 $ i & min(n, k1)
att (1) nn-
Emin(nuk) 3
Σ miin(n,))-1 3d 0
atta
Choose the value matrix as,
T3x3
03x
03x45
w
0asx3
Osx
035x45
8xx
TsxS
0sx45
0 22xx3
0sx
025
The output of the attention layer (with the residual connection) is,
x = Zn
01x2
S Ex
01228
Vn 0 xx22 ]"
min(n.k)
where, Vn
Σ
att Cx.
min(k,n-1/
Zn
Σ
att attm,n-i3'
It is easy to verify that Z, = 3 /5 for n k + 1, and that Z, < 3*/5 /5 for < + 1. This distinction will be useful later, as the value of Zn acts as signal indicating whether n > 1 or < k 1. In particular, this allows the subsequent layer to skip attention computations for indices Ti-k not well defined.
i & k; where the condition Fn -j-1, Tn-k41
B.1 Three Layer MLP followed by the Second Attention Layer
(o(") (1)
mmlp
ReLU
Norm
W-(0 -00
mip
mip
ReLU
Norm
w (5) (3)
mlp + (3)
mlp
ReLU
Norm
Figure 3: The MIP Architecture
The architecture of the MLP is illustrated in Fig. 3. We now explain how the first layer of the MLP
extracts the symbol e 2 from (1)) To achieve this, the first layer of the MLP is designed as follows:

---

(3+5)x(3+35)
(3+5)x(S))
(3+5)) (25)
w
(x++335)
Asxs
0x22
mlp
0 (x(++35)
0sxs
0g x22
(3+33)
0asxS
035x28
mlp
(1)
0 (3x(3+5)
(34-1-1)
-ixs
01xs
01x28
2-Cain
where,
effectively isolates the component associated with the symbol at position h in the sequence. This, in turn, allows us identify the value of e = for given Ixz- Before demonstrating this explicitly, we first state useful fact: in triadic sum representation, the coefficient corresponding to position n & is always greater than the sum of all coefficients at earlier positions. To illustrate this, consider the scenario in which several consecutive symbols, specifically those from positions Observe that:
Note that computing, nnmlp LN (ReLU (w!l) x 1
mlp
A + 1, are identical.
3k-1
att
Catta
36-1
3k-1 -
a
Catta
2. Catta
Hence, att,mnkk>
attn,n-i
mlp x," + Dmip computes the dot product of each symbol with v, and subtracts The operation ((1))
2.Carta from each coefficient. This ensures that only the 4thh symbol retains positive coefficient, while the rest become negative. Passing this result through a ReLU function climinates the negative coefficients, leaving only the one corresponding the n Aahh symbol. Normalizing the output sets the coefficient of the hthh symbol to one. Concretely:
gk-1
& alp = (1) ())
nip
+ b mlp
[1x(3+5)
Gix
01xs
0 01x2s] ]"
LN [ReLU(X ( kmap))
[0x(3++))
čix
01xs
01x2s| T
Where, č
[(c1
gh-1-1
ll02
gh-1 2.Cmm
1(cs
20
2.C
Note that in è. exactly one entry is equal to s-specifically, the entry associated with the symbol at position all other entries are zero. After applying the skip connection, the final result becomes:
(1) Xmm
(1) + LN (ReLU(X
Zn
01x2
e
Ciis
01xs
Vn
01x2s
1"
The second layer of the iLP is designed to compute u, from Npmppp It does so by extracting eS and Vn, and leveraging the binary mask č together with the vocabulary embeddings stored the second layer to produce up Hence, we define the second layer follows:
(1)
25

---

xx
x
03 S
x
03x
0
Sx3
s
0sx
S
S
sxxs
5xx
Sxx
0xxS
0 S
Sx
S
0xx
0sxS
w
mlp
0sx3
Isxs
AASx
S
3 Isxs
0sxS
0sxs
0x3
Sxx
0 S
Sxx
S
08x8
Osx
Sx3
xx
0sx
0sxs
x5x
0sxS
0sx5
5x3
0sx
0sx
0sxs
05x
0sx5
0sx
(2)
mlp
(3+65)x1
Hence, on computing the output, we obtain:
(22)
(())
baip
p.mip
mip mippp
01x2
01xs
01xs
un 01xs
01x2s J
3
k-1
Where, U,
Σ
att
S
Cattn
Fn-t
xm-t
To clarify how u. is constructed, we show that it can be obtained directly from Vas e and - The key observation is that é can be recovered using the binary mask è as follows, e T C where Asx is the vocabulary embedding matrix. This holds because Č is a one-hot vector that selects the embedding corresponding to the symbol at position n k. Using this, the final representation u, is computed as:
3
Un
E2g + +vn,
Cattn
Cattn
1
Cz
k
3
c
c
3
Catta
Catta
Catta
i=1
3
Catta
t=U
Hence, passing this through the normalization and with the skip connections get:
(2)
x(11) + LN (ReLU (weep (2) (11)
X.mp
(Waip"n.mp + Dmip)
T
Zn
0
012
Ex
On
Vn
01x28
Finally, on defining the third layer as follows:
03x3
5xS
0ax48
W (3)
0g5x3
0sxs
0a5x4s
mlp
xxx
Isxs
0sx18
025x3
0sxs
025x15
b mlp
0
(3+65)x1
Passing the input to the layer through the non-linearities along with the skip connections, we get:
(3)
(2)
+ LN ReLU (2(2)
(1)) (1))
b
n.alp
n.mIp
alp n.mpp
T
Zm
0 x2
S cx
01xs
un
Vn
O
x
26

---

Layer 2. In this layer, all relative position encodings are set to 0, and we use the following query and key matrices for the attention mechanism in the second layer:
W Q
K
xx (2-+33355)
Is
0
0
W K
K
Os (2++4455) +5)
Is
With these choices in place, we obtain:
(WX ((33 x.mp)) (w (3) k.mppp) K- Z,Zn +
R. (v.. un
K Z,Zn + 2 - n
V;
Un
| [ull2
u
Applying the softmax function, we obtain:
exp (kZ,Zn + 2x |T.12
la
ad1 n,
D ーj-0 exp ( exp(x
+ 2n 17,2
We can see that for all k++ < n, v; u, and only the local context around position matches that of position i.e. {z,-1-j æn for all 0, ... kj. Moreover, for any & + < n, if V; स Un then the normalized vectors are separated by non-negligible margin:
Vi
un
1
un.
34'
Note that while this gap is small, strictly non-zero. Recall that Z, gh+1 for > + 1, and Zi gh otherwise. As result, the attention mechanism favors positions such that v; u, and
i le 1. In the limit as - oo, this results in the following attention pattern:
att Unif(In),
where In denotes the set all indices E (I+ 1,...,n) n) whose local context matches that of position n, formally defined as:
I, = fie [k....,n]] | = ズn-j for all j O....
Moreover, (I.) denotes the uniform distribution over the indices in I,; that is, each position i I,, receives equal attention weight, and all other positions receive zero. We now choose the value projection matrix in the second attention layer as:
(2(2)
0sx3
I9x5
With this choice, the output of the second attention layer becomes;
((2)
m.mpp +
att
n,
1=0
27

---

Xn.alp
In
For the output linear layer, choose:
W
(SS(5)+3)
Isxs]
ba 0.
This yields the output logits:
1 zi-j In-j+1) -ー
logit,
Ozz
In
(V < k. ze-j In-s+1)
iEIn
Finally, by adjusting the notation and setting obtain:
T to reflect prediction at the final time step,
E T(V h. Fn-i IT-i+1) Em Ti Fn Fn-i
(2r+))
xT+1+1)
which precisely computes the conditional k-gram probability.
B.2 Two Layer MLP followed by the Second Attention Layer
w. (4)
A (3) Amip
Norm
Wrl 222)) + Onip
Figure 4: The MLP Architecture
We restate the output of the first attention layer:
x())
Zm
0ix2
01x2s
Vn
0
01x25 }7.
min(n.k)
where, Vn
Σ
all
all
min(k.n-1}
Z, Σ
att attm,a-;3 3
The architecture of the MLP is illustrated in Fig. 4. We now explain how the first layer of the MLP
extracts the symbol e c from (1)) To achieve this, the first layer of the MLP is designed follows:
28

---

(3+5)x(3+35)
(3+5)x(S))
(3+5)) (25)
w ())
(x++335)
Asxs
0x22
mlp
0 (x(++35)
0sxs
0g x22
(3+33)
0asxS
035x28
mlp
(1)
0 (3x(3+5)
(3*-1-1)
-ixs
01xs
01x28
2-Cain
where,
effectively isolates the component associated with the symbol at position h in the sequence. This, in turn, allows us identify the value of e = for given Ixz- Before demonstrating this explicitly, we first state useful fact: in triadic sum representation, the coefficient corresponding to position n & is always greater than the sum of all coefficients at earlier positions. To illustrate this, consider the scenario in which several consecutive symbols, specifically those from positions Observe that:
Note that computing, nnmlp LN (ReLU (w!l) x 1
mlp
A + 1, are identical.
3k-1
att
Catta
36-1
3k-1 -
a
Catta
2. Catta
Hence, att,mnkk>
attn,n-i
mlp x," + Dmip computes the dot product of each symbol with v, and subtracts The operation ((1))
2.Carta from each coefficient. This ensures that only the hthh symbol retains positive coefficient, while the rest become negative. Passing this result through a ReLU function climinates the negative coefficients, leaving only the one corresponding the n Aahh symbol. Normalizing the output sets the coefficient of the hthh symbol to one. Concretely:
gk-1
& alp = (1) ())
nip
+ b mlp
[1x(3+5)
Gix
01xs
0 01x2s] ]"
LN [ReLU(X ( kmap))
[0x((++))
čix
01xs
01x2s| T
Where, č
[(c1
gh-1-1
ll02
gh-1 2.Cmm
1(cs
20
2.C
Note that in è. exactly one entry is equal to s-specifically, the entry associated with the symbol at position all other entries are zero. After applying the skip connection, the final result becomes:
(1) Xmm
(1) + LN (ReLU(X
Zn
01x2
e
Ciis
01xs
Vn
01x2s
1"
The second layer of the iLP is designed to compute u, from Npmppp It does so by extracting eS and Vn, and leveraging the binary mask č together with the vocabulary embeddings stored the second layer to produce up Hence, we define the second layer follows:
(1)
29

---

Taxa
0
0
x 8
x
03x
03x
xx3
Isxs
s
0sxs
0sxs
0sx
sxx
Isxs
3
AASxs
0sxs
3 Isxs
0sxs
0sx
w (2)
mlp
xxx
S
0sxs
0sx
0sxS
0sx
8x3
S
5
0
IsxS
0xx
5xx
0sx3
05x
0sxs
0sx
0sx
0sx5
0xx
0sxS
0sx
OsxS
0sxs
0sxS
0sx
b (2) mip
63+
Hence, on computing the output, we obtain:
x(2)
(: (2)22((1 (1))
(2)
n,mlp
mlp n.mlp
mlp
01x2s |" Vn
= Zn
01x2
Un
0ixs
k-1
3
Where, U,
xn-i
all alla
Catta
To clarify how u, is constructed, we show that it can be obtained directly from Vns e. and "n-i" The key observation is that eie e=__ can be recovered using the binary mask as follows, & A where Asxa is the vocabulary embedding matrix. This holds because ē is a one-hot vector that selects the embedding corresponding to the symbol at position n - k. Using this, the final representation u, is computed as:
1
3
Un
Czg + 3y,
Czn-N
Catin
Catta
1
h
3k
Cavin
En
+
Catta
Bn-i
Catin
2n-k
k-1
e
Cattn
Layer
In this layer, all the relative position encodings are set and instead,
(
=
0
5x
(2+33))
Isx
0
O
0
0
(2+45)
1sx
0
With these choices, and adding a layer normalization after the output of the key linear layer and the query linear layer
LN(W
K x (2(2) LN(W (2) Q Xnmiip)
2(V un)
[ill
|unl|2
V:
un
2
u,
Taking the softmax with temperature parameter K,
30

---

exp (2x
R
[vila
vi
Un Im
att
=j= exp (2k: ki
Ivilla
Note that although this gap is small, it remains non-zero. In particular, as 50, the resulting attention pattern approaches
att (2) Unif(In Ix),
where Zn denotes the set of indices i E (k...., n) whose local context matches that of position n. Formally,
In-j for all 0, k). I, = : (k....., nn) | æi-1-j
In contrast, Zg captures earlier positions E {O...., k 1} whose (normalized) key vectors are aligned with the normalized query position n::
ix = ie
1}
Vi
k,
|u,
On choosing,
0sx3
1sxs
0
We obtain,
n
x(2
(2))
(33
att
S
[nal + [Tn
sell., wil
For the output linear layer, we choose:
W OSx(5S+3)
Isxs]
b, 0.
which results in,
[[V1 & マ k, Iz = kn-j+1)
logit,,
S I(VI & ブ k, Zy-j xn-j+1) + It
[n,
[Ixl
iEI, oft
< It Fn-j+1)
ーーー M < k, x2 xn-j+1) It E
Te
Now, by substituting n = T. and aiming to predict the next symbol given the current context, we obtain:
T b I(A0 Ei ki Fn-i = FT-
D =
Σ
I[VI & & k, In-i = DT-++1) + Ig
31

---

Exet, (T+1,
Ex= I(1 Vivliis k. on-i = xr- + Ix
This corresponds biased conditional k-gram model, where the bias originates from the second term in the expression above. Notably, this bias vanishes in the limit as T - 0c. The reason is that the numerator of the second term remains constant with respect to sequence length, since Ia is fixed and does not scale with T. In contrast, the denominator, given by Σ I(V 1 & k, In-i Fr-i+1), grows with T, assuming the underlying Markov chain is irreducible (i.e., every state is reachable from every other state). Consequently, under this assumption of irreducibility, the influence of the second term diminishes as T becomes large, yielding the approximation for a sufficiently large T):
Σ
_n I[ & i & k; In-i = FT-i+1)
(ogity(xr+1)
'T m
& i k. Fn-i = xi+1+1)
32

---

C Two-Layer Single-Head Gradient Analysis
As discussed in Sec. 5, we analyze simplified model that retains only the key components of the two-layer architecture described in Thm. 3-specificallly, the positional encodings and the attention layer-whie treating the remaining components approximately optimal. This reduction enables tractable yet faithful reproduction of the in- in-context learning (ICL) behavior exhibited by the full model. Our approach is closely related the reparameterization techniques used by Makkuva et al. [18] and Nichani et al. [20] in the context of gradient flow and gradient descent analysis. Alternatively, can be viewed as a form of optimal parameter initialization, in the spirit of Edelman et al. [10] and Nichani et al. [20]. We reiterate the assumptions and provide additional details whenever necessary.
Data Assumptions. Following Nichani et al. [20], we assume prior distribution (=(() over irreducible and aperiodic Markov transition kernels т, such that:
Assumption 1. (Assumption on P:)): We consider prior distribution transition matrices of first-order aperiodic and irreducible Markov chains such that there exists a constant y > for which, transition kernels drawn from (.).)), the following conditions are satisfied:
1. Minimum transition probability:
min i(s' 8.8"
7²
2. toiviaivial mixing: E -xx
3. Permutation invariance: For any permutation a on [S], To 4
4. Uniform expected transition matrix: [.[[r]
1sl
Model Simplifications We simplify the model as follows. More details can be found in App. C.I.
The positional embeddings used for the keys the first attention layer are scalar-valued, i.e. (1.K), x Pa. where is a scalar.
(ii) The query, key and value matrices (W(1) Q.K.v) and positional encodings for the value (p (1,)) p,n-; exp(p:)/ (==((x)pp;;)
for the first layer, are chosen so that the attention weight for position i is given by
att (1)
{0, 1), and the corresponding value output is
E att, (1) n,,-
e P.-1"
(iii) We assume the MLP is optimal-i.e., it outputs the correct u, even from suboptimal v.s this holds for first-order Markov chains via skip connections.
(iv) The query and key matrices for the second layer (2(2)) Q.K) are chosen such that the attention
in this layer is defined as att. = exp (ag- (r)) / (2 exp(a2 (v,, (v,, un>)). where
02 E R is the attention scalar.
(v) W is chosen so that the final logit logity LT att (2) n.i Ex, RS
Loss function. We minimize the cross-entropy loss for next-token prediction [20]:
((0) - ET-Ps. Fo
Σ (xT+1 fo.r) log (logit,(s) + +e) SEES
(2)
where e is a small additive constant, (-) denotes the prior over Markov transition kernels, and (s) is the logit corresponding the symbol
Training algorithm. We denote the set of trainable parameters as (p. 02), where p [Po, ... ,PT]" "Pr]" represents the positional scalars and is the attention scalar. Following the approach of [20], we employ a two-stage training procedure. (i) In the first stage, we optimize positional scalars p using gradient descent for Ti steps with learning rate 7)1 (ii) in the second
33

---

stage, we freeze and train only d2 for additional T2 steps using distinct learning rate 122. A key difference in our setup lies the treatment of layer normalization: during the first stage, all layer norms in the MIP are omitted, whereas in the second stage, they are reintroduced. This stage-wise strategy facilitates more tractable theoretical analysis. The training algorithm is:
Algorithm 1. (Training Algorithm).
Input: learning rates 7) 1/2; steps Ti, T2; small scalar 02,0
Initialize (0) 0 for all i E {O...., n2); 02 02,05
Freeze all other parameters of the network
Stage 1: Train p; without norms added to the MLL for t 1...., Tj do
(() (t-1) - - Vp((6(-1)
(0) g(t) (p) 1r 02 (2.0)
end for
Stage 2: Train d2 with norms added to the MLP
for 1 = Ti + (t-1)
Ti + T2 do
a2 (t)
7/2 Vay L(f(t-1))
02
g(t) {p:
(1)) |o 02
end for
((T++3)) +T²))
Output:
Under the assumptions outlined above, we restate the main theorem presented in Section In addition, we introduce new theorem that addresses the generalization properties of the model in the inductive learning setting.
Theorem (Convergence of the training Algorithm). Suppose that Assumptions (i)- (vi) hold and that the Markov sequence length I' satisfies T + > poly(y-1 (S)) for some constant 7 O. Then there exist 0, learning rates 7ı, 772, and step counts Ti, T2, such that the output of the above two-stage training algorithm, 0 (1P;) â2). satisfies
c(0) L* z log 1) (T + 1)
for a constant c independent of (7, and the optimal loss L E (1/S) T (ss 8) log (s'' $)}.
Theorem (Inductive Generalization): Let मे be transition matrix satisfying min,, #(s') s) >> v/S, and let 9 denote the trained model from Thm. 7. Consider sequence 80, generated according स. Then the model satisfies:
log(n 1) (n + 1)c7
sup fs ....s.................. sn). 7(s') sn) s
This section organized follows. We begin by presenting the simplified construction in App. C.1. Next, we demonstrate that during the first stage of training C.2, the model converges to the optimal attention map for the first snyer-ssecifiicall is optimal under a first-order Markov chain. In the second stage of training C.3, the model sharpens its attention, ultimately converging asymptotically
learns attend the previous token, which
a near-deterministic distribution over desired token. Finally, we provide arguments for the result in the context of inductive generalization, utilizing results from Nichani al. [20]. We would like note that our analysis for first-order processes similar to that of Nichani et al. [20], but the architecture is fundamentally different.
34

---

C.1 Model Simplification
We adopt a www-parameter regime by isolating the core components the Transformer architecture that are essential to our analysis, while keeping all remaining parameters fixed at their optimal values. Here, optimality is defined with respect to first order Markov chain setting. Under this setup, proceed by stating the following assumptions.
Emb(x,,)
[01x3
01xss]
TR
In the first layer, define the positional encodings corresponding to the key as follows, while assuming that all other positional encodings remain fixed and set to zero.
().K P
{p:
1
0 01x(1+65)*', if {0, 1,,1,2.... .. ,n),
Layer 1. We paramatrize the query and key matrice follows:
w
01x1 O (2+00))x
1ix2
0 x1x68
K
(2+68)xx
0 (2+65))((+055)
(1))
O1x3
1 1ix
0 x55
0 (2+65)xx3
(2+65)
0 (2+65)
This then leads to:
(())) (Emb(xn-1) + p K, w Emb((xn) pi K
Hence, on passing through the softmax function:
exp(p;)
att (1) n E1= exp(pt)
Choosing the value matrix as,
Taxs
Osx
03x4S
(((1)
Ossx3
Osxs
038x48
0sx3
Isxs
0sx45
028x3
Osxs
028
Hence, can obtain,
(1)
= 01x3
S
er.
01x2s
Vn
01x29 1
12 >
where, Vn att
1=
MLP Layer 1.
The first layer of the MLP is defined as:
(3+335 (xx33)))
(3+
(2S)
(1)
0xx
0sx28
(3+38)
0sx
0 0 25
(3++33)
0sSxS
038x29
b (1)
0 (3xSS)))
01xs
01xs
O 01x28]
mlp
35

---

Hence, after applying the skip connections, we obtain:
(1) x nlp = x(1) + LN RRLU (w(Uxgh) (1)
= 01x3 | ex, 1ixs
0ixs
vn
01x2s ]"
MLP- Layer 2.
The second layer of the MLP defined as:
03x3
O
0
3x
3xS
xx
ssx
S
0sxs
0sx
0sxs
Ssx
0sxs
sxs
0xx
0sx
0sxS
0sxs
0sx
0sxs
W:
0sx3
1sxS
0sxs
0sxs
0sxs
0sx
0sxs
mlp
Sx3
sxs
0sx8
Osxs
0sxs
0sxs
0sxS
Sx3
S
0sxs
Osx
0sxS
Osx
0sxs
Sx3
Sxx
0sxs
Osxs
0sxs
0sxS
0sxs
2
mlp
68+
Hence, on computing the output, obtain:
(2)
(1)
LN
ReLUU
(wlel (1)
S + b(2)
p.mpp
01x3 es
0ixs
Un
Vn
0 01225
]
Where, Un
Note that, in this setting, the optimality condition well defined due the first order Markov property. In particular, un can be recovered directly via the skip connections, as it depends only on the current token.
MLP- Layer 3.
The third layer of the MI.P defined as:
0xx3
Sxx
03x1s
0a5x3
Osxs
0asx
W
mlp
0sx3
IsxS
0sx4s
02
0sxs
025x48
b alp
0
x1+
Passing the input to the layer through the non non-linearities along with the skip connections, we get:
(3)
(2)
LN (ReLU
((1) mipg
+ b
n,mlp
xn
T
1x
01xs
u.
Vn
Iv..a
01xs
Layer In this layer, all the relative position encodings are set as and instead,
W (2)
a2
O (2-35)
Tsx
0
W(22
0
0
(xx(2445)
Isxs
0
36

---

With these choices, and adding a normalization after the output of the key linear layer
W (2) (122)((3) X n.lpp
xm.alp/ 02 (Ve: Un
Taking the softmax, we obtain,
exp(u2 (Va, U,))
att (2)
-5- exp(a2 (vj, un))
Finally the logit is computed follows (selecting the optimal value and projection matrix):
logit,
> 1=0
(2) all -xa
Motivated by this parameterization, we construct a simplified model that abstracts away the complex transformations occurring within the intermediate layers of the network. During stage of training, we remove layer normalization; accordingly, we omit layer norms from the simplified model as of first
well. These can be reintroduced later during stage two of training. However, in the order Markov chain, layer normalization does not affect the model's behavior, since both u, and v, converge to one-hor embedding vectors. With these considerations, the model can be concretely reduced as follows. Let Zn E R n=++
where Zi = (Bx,, Ui). Morever, let:
X, [Ezo
Cxi
Exs
Ez,]
Hence,
logit d, =e² [X,, softmax (n)
Note that the logit computed above corresponds to the case where In notation to account for causal masking, we obtain the following:
Using slightly different
Po
-
00
-00
-
00
00
P1
Po
-00
-00
A0)
Pn
Pn-1
Pn-1
Po
(2) 02 T
Hence, logit,, 2 [N, softmax (a2
(ofttmax (4(1) X² (22) exn
X, softmax
(X.:1)
We now use this simplified model for the gradient analysis.
C.2 Stage Analysis
We use the simplified model described in the previous section. We begin by quantifying the loss as follows:
37

---

I(0) 0
P(r)
Σ
π (8n+1 = 8115) log (logitss. te)
Let I, [1,1,". 11" E R Hence, we can then define:
logit 2/m e [X,, softmax (2))
Assuming, that 0220 0220 is very small, we can first order taylor approximation. We also assume that En = s:
1n + 0220
In
logit,, sn
2 E
X
1 1 n
Zn
n
1n
T X Zm
T X In Eg
1n T
+ 02.0
1
11
1
1
d2,0
ex Zn
c 2, 1n
1n Tnn
n
n
n
Note that, on computing the quantities separately:
T en,
(b)
1 T Zn
{ vi)
i=0
exp(p,)
Σ (ean
1=0 Ei-o exp(py) ex.,)
exp(p;)
(Bam ex--s)
= 0 exp(p))
n Zn Σ
exp(p,)
(Bon Cz Cxi,, Xez, ea)
(- E/=o exp(p))
Hence, on substituting the above quantities:
logity,, (5)) +
$2.0
exp(p;)
{ee, e exn,, (Ex, Cx, Pex
1=0 En -0 exp(p,)
02,0 (8) + Σj exp(p;)
exp(pı)
{Iz----s 1zy-s 1 1xp--- Pi (5')3
n
Hence, computing the derivatives of Pm. keeping all the other variables constant, we obtain
dlog (logit
dL(O)
Σ
(sn+1 = s'|s)
H- P(x),xx
dpm.
dpm
38

---

Note that for sufficiently long sequences, px,((s') 2 pe (s'). Thus:
dog (logils..n)
0220
exp(pm)
cxp(p, 01)
qim
exp(p} a1) 9:.
dpm
8 px (s')
=0= exp(p;)
Σ
1m
j-0
Where, 9i
===-===" -,,=0'
=,=== Ms
Using the tower property of expectation,
dT((0)
02200
IE 1 7(a))
Σ
1
(5'1s) exp(pm)
dpm
5 km
(x' Σn- exp(p,)
(P(x:;
s', 3/
8) P( = 8) (=((s))
+ 02.0
Bxx-((r) P(x))
Σ
T(8'ls)
exp(p;) exp(pm)
S. knn
33m 8',8 f=u (s(
(P(x; = s'; = =8) 8)- - -P((;---- s)
We define new quantity.
Σ ()'1s)
(T.
(P(2;= 8 s) -P(xi-; = s)Hx=(())
(x(
(s'(3))
(P((x; =. ,,, 5) pr (5)) (from stationarity)
(8))
(s'|s)
P(x;; 8. =s) V (s'1s)
8. (s() pn (s())
(s')
Σ (5'(s)) Pl(xi 81.11- 8) -1 (by marginalization)
(5')
gi.j Fm))(() (g).
Hence, on substituting this back, we obtain,
dL(())
022 S kn
exp(pm) vー= exp(p;) Fim Pim( (TT)
exp(p;) exp(pm) o (
Ex-P(x)
dpm
Therefore, rather than considering the entire expression, we focus on specific row (i.c., for a given i), and obtain:
(
dT(6))
02
En-P(rr
exp(pm)
Σ
exp(p)) exp(pm)
- (TT((T))
dpm
S. kn
Eo= exp(p;)
3-0
39

---

To obtain the optimal attention map for first-order Markov chain, the goal is to demonstrate that, for every row, the gradient with respect pi greater than that of any other token and is also positive. Due to the properties of the softmax function, if this condition holds over the course of training. then the model will asymptotically attend to the previous token, thereby converging the optimal attention map. Hence, the requirement is essentially, for every row:
dL(0)
dL(O)
Vk (0. n}
dpi
dpp
dL(O)
dp1
In this section, we will show that given the prior assumptions, this is indeed true. Towards this, we denote dgm..
dL(0)
With this notation, express:
dpm
gm
02.0 S km
exp(pm)
9t.m
Σ exp(p,) 1-0 Σ 2 exp(p;)
Ii.
exp(p;)
Now, suppose we initialize all positional encoding scalars to zero as defined the training algorithm. Then, at initialization time t 0, we have:
0220 1 10-0 S
Fi.m
gi.s
Vi
i 1+1
Due to Lemma 7, we observe that at initialization time t = (when all positional encoding scalars are initialized to zero), the gradient of the attention logits is maximized in the direction corresponding to position 1, This implies that the model is initially biased to attend more strongly to the immediately preceding token, aligning with the structure of first-order Markov process. Concretely, we observe that:
-o -gr Ak (0, dg... Ovie mm]
n), Vie (k. m)
and
Hence, under the given initialization and model assumptions, performing gradient descent yields the update rule:
P1+1
Pt 7 Vp, L(0),
where
[Po.:
P1.t
pn,)] e 1B#+1
Here, Pi, denotes the i-th positional encoding scalar at iteration t. Therefore, after the first update step (i.e., at 1), we have:
P1 Ph.1=1
vk
20.
reflecting that the gradient is largest in the direction corresponding to position 1.
We need to show that this property holds for all iterates. By Lemma 3, we can show that for each iterate of the preconditioned descent, the following holds:
P1. Ph.t Vk E {0, 2..
40

---

Furthermore, by Lemma 3 and Lemma 4, we observe that Pi tends to infinity it increases monotonically over time. In contrast, the other positional encodings do not diverge, as they do not exhibit similar growth. Due to the nature of the softmax function, the rapid increase of pl suppresses the relative influence of the other components, effectively preventing their growth. In addition, Lemma 4 provides an estimate of the number of iterations required for the iterates to converge to the optimal attention configuration. Therefore, assuming a sufficiently long training period during the first stage, the model successfully recovers the ideal attention map corresponding to a second-order Markov chain:
Vi Cx
Hence, in the second stage of training, the inclusion of layer normalization does not significantly affect the outcome, provided the Markov chain is sufficiently long and the training duration adequate, and hence, we will assume that v:/ v|ll R Ca;-1 which is the optimal attention map. We note that the gradient analysis for the first stage of training closely mirrors that of the first stage in Nichani et al. [20], and is directly inspired by their approach.
C.3 Stage Analysis
We observe that, at this stage-assune assuming the attention map learned in the first stage is optimal by having the appropriate sequence length and training time-our simplified model aligns with the one presented in Nichani et al. [20]. This alignment arises because, after the first stage of training, both models converge to the same optimal attention map in the first layer and subsequently train single scalar parameter in the second attention layer. Therefore, we can directly invoke the theorems from Nichani et al. [20] to analyze the convergence behavior of our model during the second stage of training. Accordingly, we state the relevant theorems below without proofs. Using the same set of assumptions outlined in Appendix C, we restate the key lemmas and results from Nichani et al. [20] as follows.
Lemma 1 (Lemma D.8 in Nichani et al. (20]) Let (fx)
(TL) ( )io al""). where (pg (T2) yo-o denotes
the output after the first stage of training. If a2 satisfies exp(a2) & exp(a²) = CG,s n 1/12 log 1/6 n, and 02 0 (at initialization), then
1
dL(O)
1 1
6
202 0.
da²
This clearly indicates that the negative gradient remains positive, causing d2 increase indefinitely under gradient descent. As result, the attention mechanism progressively sharpens, effectively converging to a hard attention regime. Over time, this leads convergence to the optimal 1-gram estimator as the number of training iterations grows. For a thorough discussion of the time taken to converge, we refer readers to Nichani et al. [20] (see Lemma D.10). Furthermore, given the equivalence between the models, once training is aapplette-ssumminng
sufficiently long sequence and adequate training duration- they are essentially identical. Therefore, we defer the proof of convergence of our model's loss the optimal loss (Theorem 4) to the corresponding argument in Nichani et al. [20] (see D.7. Proof of Theorem 4.4). Finally, due to this model equivalence, and after the completion of both training stages, we again refer to Nichani et al. [20] (see H.4. Proof of Theorem 4.5) for the proof of Theorem 8.
41

---

D Two Heads in Layer One, One in Layer Two Perturbative Analysis
Understanding the convergence behavior of a Transformer Markov chains of arbitrary order in the prior setting (Appendix C) presents significant challenges. central difficulty lies in the dependence of accurately estimating u, on the precise recovery of v,, In earlier constructions, this estimation was enabled by complex non-linearities such as layer normalization, which makes direct gradient-based analysis intractable. To facilitate analysis, focus simplified architecture: two-layer model where the first layer includes two attention heads, as described in Section 3. Due to the intractability of analyzing the full model, we adopt perturbative approach in which only a subset of parameters is trained, while the remainder of the network is fixed at optimal values. This aligns with our goal of exploring the behavior of wwwpaaramtterzeed version of the model. Within this setup, we consider arbitrary-orde order Markov chains under the assumption that the first head is optimally initialized to recover u, Consistent with the simplification in Appendix C, we limit learning to the key positional vectors of the second head in the first layer, parameterized by single scalar. All other network components remain fixed. The second layer of attention not assumed be learnable; instead, we begin with softmax function at very low temperature, held constant during the initial training phase. This temperature is later increased to infinity. Even under these constraints, analyzing convergence of the positional encodings to their optimal values remains analytically intractable without strong assumptions about initialization, training procedures, and the data distribution. Nonetheless, we regard this work as a first step toward understanding convergence in higher-order Markov models within more realistic Transformer architectures.
Training Algorithm We introduce training algorithm that updates only the positional scalars ps for the key vectors of the second attention head in the first layer, for all i E (o...., n2.. Here, D ((--1) denotes data-dependent diagonal preconditioner.
Algorithm (Training Algorithm).
Input: learning rates 7), steps Ti,
Initialize p; (0) 0 for all i E {1..... n and Po = 00;;
(0)
Freeze all other parameters of the network
Set the softmax temperature of the second attention head to very low value.
Stage Train p; without norms added to the MLP
for ( 1,..., Ti do
(+) p (tt11) Th (D VPC(0(--11
g(1) (o. (0 )=-0)
end for
- ((T))
Output: θ
Add norms back to the MLP
Set the Set the softmax temperature of the second attention head oo
Data Assumptions We introduce the following data assumptions on the prior over Markov transition kernels, first for second-orde Markov chains and subsequently for higher-order Markov chains.
Assumption 2. (Assumption on P for Second-Order Markov Chains): We consider a prior
distribution P. over transition kernels of second-order Markov chains defined over binary
state space S = {0, 1J. Each transition kernel π ~ IPx specifies conditional probabilities (ss' $1, s2), and satisties the following assumptions along with stationarity:
1. Time Reversibility: There exists stationary distribution p(81, 82) such that (((1, 82) (T'') $1, 82) = 4(82, s') ( $2, s').
2. Transition Preference: (0 0, 0) > (00 1,0) and r(1 | 1,1) a(1 0).
3. Second-Hop Likelihood: (o. 2 s, s) > 1.
42

---

Assumption 3. (Assumption on P for kdh-Order Reversible Markov Chains): We consider prior distribution Px over transition kernels for kth-order Markov chains on finite state space S, with lifted space = S" such that each transition kernel æ Fx satisfies:
Irreducibility and Aperiodicity: The chain (X₂) is irreducible and aperiodic; 2. Reversibility: The chain is reversible with respect a stationary distribution x;
3. Lifted Representation: The chain admits lifted first-order representation with transition kernel P(a - b) Pr((X............ X = b (Xi-1..... Xx_ = a);
4. Spectral Assumptions:
Under detailed balance with stationary distribution x, the transition matrix P is self-adjoint on (A, .T, with eigenvalues 1 d > X2 2
> Av and orthonormal eigenfunctions {Øm):
5. Non-degenerate Projection: The eigenfunction projections 8m, defined by
8m
Σ Φm(a) 6m (b) (Tb f(a),
a.deAA
satisfy Bm for all m, where f : is an observable function.
Under the assumptions outlined above, we state the following theorems.
Theorem 9. (Convergence of the Training Algorithm Second Order): For all second-order Markov chains satisfying Assumption 2, there exist, learning rate 7h, and a step count Ti, such that the output of the training algorithm (Alg. 2), {pili_o> satisfies the following as Ti 00:
exp(p1)
exp(p2)
Σ exp(p:)
= exp(p:) 2
Morcover, the softmax temperature is increased infinity, the model reduces to conditional 2-gram estimator.
Theorem 10. (Convergence of the Training Algorithm -oddrr):):: For all -odde Markov chains satisfying Assumption 3, there exist a learning rate and step count Ti, such that the output of the training algorithm (Alg. 2), A = {p satisfies the following as Tj oo:
exp(p1)
exp(p2)
exp(pk)
12 2Í cxp(p;)
En exp(p:)
2:-0 exp(p:)
x K
Moreover, as the softmax temperature is increased to infinity, the model reduces to conditional k-gram estimator.
We begin by describing the model simplification and subsequently present the convergence analysis.
D.1
Construction
Note that we assume one of the heads (u,, to be optimal, we only need to focus on the other head. Since, we are interested understanding the perturbative analysis of the second head with regards to the first, we define the layers follows:
Emb(
[
01x5s]"
6553
In the first layer, for both the heads use the following relative value embeddings
43

---

(heed:1).V
(hed::2). P
0
[1 0]²
for i < k
P
k.
Layer 1, Head 1. We assume that the first head is optimal and hence, have an architecture that outputs the optimal attention map. Specifically for a k.hh -order markov chain, we get the following attention score:
U, (ex, + can, +... + Exn mi1)
Hence, the output of the attention head is:
(head:1) n
=
[
01x2
01xs
01x2s
01x2s ]",
Layer 1, Head 2.
We will use the same relative position embeddings, particular,
[bead:2),K
{Po
[ 1 0 01x(1+6))'' if E 1,2,,,.... ... ,,), ,n),
We paramatrize the query and key matrices such that:
W (head:2) K
whbaad: (Emb(z (n) () (head::2) "Emb(xx)) ) Pi
Hence, on passing it through the softmax function:
cxp(p:)
2 Σn= exp(p:)
(1)
Hence, once we keep the value matrix fixed (not learnable):
Tax3
03x
03x45
OsSx3
Osxs
038x48
Dsx
Tsxs
05x45
0 025x3
0sx
028x48
The output of the second attention head
(head:3) :2
0
01x2
0xxs
01x2s Vn
01x2s F
min(n.))
where, Vn
Σ
att Cza-t
5
Therefore, with skip connections:
x (1)
xx2)
+ X,,
(1haad:) (head::)
X
0 1x2
e u 0
V,
01x2s ]"
44

---

MLP. The layers of the MLP are defined as follows:
The first layer for the Ms.iss:
(3+3))(3+35)
(3+-5)
(3 x
(1))
(3+338)
ssxs
xx22
5x(3335)
Tsxs
Sx2S
33+(3335)
0
35x28
b mlp
(3+68)x1
The second layer for the Ms.
(3++3)
(3+ (5)
(3+55)
(x+399))
SxS
5x22
W (2)
(x(3335)
0sxS
5 xx
(3+38)
IsxS
23x(33355)
028x8
28x29
(2)
(3+65)x1
Hence, passing this through the normalization along with the skip connections, we obtain:
T
n.plp
0
01x2
y1,
U.,,
Vx2
x
Layer In this layer, all the relative position encodings are set as and instead, and paramaterize W (x(2) such that
K
2(vi, un>
(W
x22 (W (2) imll) Q (mmmpp))
||v:|12
|unll
V;
U1,
2
vall2
u,
Hence, the attention is computed as follows, using softmax with temperature parameter s:
uu vil u.,
exp
all
Ei= cxp(1i
uv. 1
Hence, obtain,
logit, Σ
(2) att -Ex;
D.2 Convergence Analysis
We can effectively reduce the model described in App. D.1 to simplified form, analogous to that in App. C.I. Specifically, in this case, obtain:
Let 7,, E R1+11 where
2, = ( (Ez, + Ean-s + (en-2) 1
45

---

Moreover, define:
X, = [ezo
6x1
6x2
6zn
Then,
logit, " 1n = e
[Xn softmax (kZn))
Using a slightly different notation, we can express the same as:
Po
P1 - 00
-00 00
A
Pn Pa-1
Pn-2
po
1 R.
logit
eg [X, softmax (kZzz
X softmax
k (.-sotmax(A(1)))
Exn-t
k
A 1 (X..k))
Clearly, the above expression closely mirrors the reduced model presented Nichani et al. [20]. Since have picked the softmax parameter K to be very small, 1 approximation to simplify the expression. Let 1 = [1,1, ,1] e R=+ We also assume that
apply first-order Taylor
Tn 81, Fn-1
82,
Fn-k k 8k-
ln + tk.
In
logit d
2 Eg
X,
1,
Zm
n
N In KR.
exx,n 7n
X, n 1
1 7m
12
x In
T E
X 1n
T E Xx 7m
1 7m
n
Now, compute each term separately:
nA (5)
n
1 En = Σ
Exp 11 Vi
k
Σ
1
R-
Σ Exn-11
Σ
exp(p;)
k
1
2,=0 1 exp(p;)
0-0
exp(p,)
Σ
6an Peins
ーーー exp(p;)
i-j-0
exp(p,)
X, 7m
Ca. Czins
(ez,,
1=0 1 ύ-0 exp(p,)
Hence, substituting the expressions above.
exp(p,)
logita, n (xx,)(s)) + kn
i=0 Σn exp(p;)
[=0
46

---

Co Ex/) - (exn-l: (z._))
*
exp(pi)
(s') +
kn
z= exp(pj)
(1 == 1x==0'
la
M (6)))
==
Hence, by computing the derivatives with respect to Pm, while keeping all other variables constant, we obtain:
d dlog (logity..
dL(0)
B(n)(()) on
Σ 11
T($n+1 = ss $1 s1...., st)
dpm
S/
dpm
Note that for sufficiently long sequences, (s') 2 px((s)) Thus:
d log (logit....
exp(pm)
exp(pj)
dpm
Skk en p (s')
1-1 2j exp(ps)
Vim
1= Ej exp(pj) Widd
where
Pisi 1 1m---si
1 mm
1 1m----oi Ms((')
Using the tower property of expectation,
dL())
K |S) kn
E B.))(r)
Σ
((f((11.....55)
exp(pm)
dpm
Ma(s')
Σj= exp(p,)
ism
(1(x;; s', 31) P(zi-m s1) (x())
T(s' T(6'(s1,...,33)
exp(p;) exp(pm) Eo exp(p;)
E Ex
| kn
(s(
13m 1=1 s', j=0
(P(x, ='' si) PP = st) pT
Hence, instead of looking this in totality, for particular row (i.e. for particular i, we obtain:)
dT((O)
K
E
P(rr)))
Σ Σ
k
T(8'[11,.... 88)
exp(pm)
dpm
S kn
(x))
Σ Ej exp(p;)
I=ls'ssi.....
(P(x = 5', 3. = 31) P( = s1) un(('))
T(s'|s $1.
ss)
exp(p;) exp(pm)
+ SSk 1 kn
Bx-P(r)
Ma
Ej explp;))
1=
118..81.....44 1-0
(P(2x 8', 5' si) -P(xi si)
(4)
47

---

We define new quantity,
iTs s1, .8k)
Σ
(P((I; 5 = si) - P(xi_ç = st)
Mx(8')
T T(''|s1,-..,, ,8k) (P(x; = 1i = si) - pr (s())(((ss)) (from stationarity)
Σ
Ma(s")
s1 8k)
a(s'[s1,... sk)
Σ
Pp(;; 8 ,i- s1) Σ
p (=()))
(x(s')
T T(8'/s1,... 8k
Pr(;; = ,'
-z-j = si) - |S| R-1
(by marginalization)
Ms(s')
8',81, 8k
Moroever, we also define:
91.j
Ii.dd
Hence, on substituting this back the original equation:
dL(O) dpm.
6 SS kn
exp(pm)
exp(p;) exp(pm) Ej explp;)
P(x)(()
T2 mi
29:3
- exp(p;)
k
D.2.1 Warmup: Second Order Markov Chain
We begin by presenting the dynamics for second-order Markov chain as warm-up. Since the vocabulary binary, we have [S] Therefore,
k Σ
giid
Σ
T(s' $1)
Pp = 5', =
s1) ATs' 82) mx(0")
P(x: 5 s1, =s2)
(8")
8'.81.82
Enumerating all possible values of $1, 82 E (0, 1}, we get:
TTs' 10,0)
Me (s') P(x; s', a (, 0)
P(z, = 5 = =0)
9t.
(x()
TTs' 1,0)
( 1,0)
P(x, -,, 1)
P(u: s S 0)
(n((8)
Hn (s))
TTs' 0, 1)
i 0, 1)
Pl(x, ,,, =0)
P(x s 0))
(n((8))
Nn (s')
TTs' 11, P(a 8, 1)
( 1,1)
P(: 5 , 11) -
P
(+(8²)
(s()))
Using marginalization, simplify terms such as:
48

---

(ss' 11,0)
P( 8 1) (Ts' 1,0)
P(xr4 8', 01
(x(s')
(x
P(x, 8' 1) + P( 8', 24415 0)
Σ
7(s' 11,0).
Pz (s
(a(s')
Σ
(<ss 11,0)
(by marginalization)
Mr (s((s))
Σ
T(s' 1,,0)
Similarly, obtain:
Σ (s' 00.1) P(x; = 5', 0) + (xx(5)
(T'' 10.1)
P(x, = - =1) =
Mn (sn
Therefore, substituting these identities back into the expression for gi.j. we obtain:
Jij = 2. Σ (ss 0.0) (5))
P(z. 8, 0) + Σ (s' 1.1) fx (8')
P(x; = = -1)
In an ideal second-order Markov chain scenario, we expect the positional scalars and p2 to approach infinity, while the remaining positional weights remain relatively small, This reflects ideal attention behavior: the model should focus primarily on the two preceding tokens, consistent with the structure needed for accurate k-gram (specifically, bigram) estimator. Hence, the requirement essentially:
dL(O)
dL(O)
aL(O)
VkE (0.3....,n) ,n
dp1
dp2
dpk
We denote dgm,,
dL(#)
With this notation, we can express:
dp...
exp(pm)
exp(p;) Li exp(p;)
0220
dg
Jim
Gi.i
S
kn
Σ exp(p;)
According to assumptions (Assump. 2), the chain satisfies detailed balance. Additionally, by Lemma. 5, we know that the values It decrease monotonically:
9..0 91.1 91.2
Fi.
Using this monotonicity, we can bound the weighted average:
exp(p;)
exp(p;)
Σjz exp(p;) Di.j &
Σ exp(p;) -up St.
< sup gi.s = I..0
(by monotonicity)
49

---

Now, suppose we initialize all positional encoding scalars to zero except for -0. Then, at initialization time have:
dgm,,,1=0= =
22
Σ i
E
| kn
i 2-1
t0 =
Since we assume detailed balance (Assump. 2) we can now deduce:
TT'' ((,0) PP(:; 8, 0) Σ T 11,1)
PP(2; - 8 1) -111
gig
2
(5()
Mx((')
2
Σ P( 00 00 =ss (Ts' 00,0) + Σ P(zi-s = 1 x, = 5') T(S' 1, 1)
Σ pl |s, (1))) (by time-homogeneity and detailed balance)
= 2.
Under the assumptions on the Markov chain (Assump. 2), observe that at initialization time t (when all positional encoding scalars are initialized to zero, except for -00), the monotonicity property 9t,1 > 9i.2 > 9i,n holds.
Moreover, by assumption (Assump. 2), we know Jú,1 > 1, which implies:
don.00 > -dg...0 ................................ dgn.00 for all
This contradicts the ideal case described earlier, where we would prefer:
m0000 = =ddg,..0 > -gn.
AK nI,, Vie mk....,,,
(5)
and
0g 2 0
Towards this, we first define a row-summed quantity:
dg
dam.it
To satisfy the condition in the ideal case (Eq. 5), we now introduce preconditionee--bsee optimizer:
P++1 P - m) D V V.. (0)
where
P: [po.
Pı.t
pr,1]' R
D. 1 diag
0, doud
1, I.I.....1 e ((in-)x((+1)
dyz.t"
Here, Pi,t denotes the ith positional encoding scalar at iteration f. Therefore, after the first update 0), we obtain:
step (i.e.,
P1.1 P2 Pk. Vk 3..,....,7 n)
......7..
50

---

Note that while we assumed = -C for analysis, in practice we can set it to large negative constant ensure that i. > We need to show that this property holds for all iterates. By Lemma 3, we can show that for each iterate of the preconditioned descent, the following holds:
Pt, P2, > Pk.t Ak 30.3.....1
Furthermore, by Lemma and Lemma 4, we observe that pi and p2 tend infinity, as they increase monotonically over time, whereas the remaining positional scalars grow at a much slower rate and thus either diverge to -00 or converge approximate time required for the iterates to reach the optimal attention configuration. Therefore, by the end of the first stage training (as Ti - o0), we recover the ideal attention map corresponding to second-order Markov chain:
finite constant. addition, Lemma characterizes the
2 (C + Ex4-2)
By reintroducing the layer normalization and taking the limit ls - co, recover the conditional 2-gram estimator, as shown in Appendix A.
D.2.2 Any-Order Markov Chain
Following up from the previous section, we follow the same philosophy as before. Similar to before, we expect that an ideal scenario, we expect the positional scalars PM to approach infinity, while the remaining positional weights remain relatively small. This reflects ideal attention behavior consistent with the structure needed for an accurate k-gram estimator. Hence, the requirement essentially:
dL(O) dp1
aL(O)
dL(O)
dL(O) dpg ii
Vqe k + ......n)
dp2
dps
We denote dgm.:
aL(0)
With this notation, we can express:
dpm
0g,0
exp(pm)
exp(ps) 2i-0 exp(p;)
dg
E F(T))
Jim
S
km
j0 exp(p;)
Now, suppose we initialize all positional encoding scalars to zero except for Po -00 Then, at initialization time
0, we have:
dgm,..t=0
[8]k kn i [E.r(r)llnmml
00
Σ 3-1 i-1 E (nwe(mllus)) Vidi
dgm.o.t0-
Under the assumptions on the Markov chain (Assump. 3) and Lemma. 6, we observe that at initialization time 0 (when all positional encoding scalars initialized to zero, except for Po = -00), the monotonicity property Jt,1 > Ji,2 > ... > fi.n holds. Moreover, by Assump. 3, we
know gt,1 > which implies:
-gs..0 d 00
10... 0, Vi
This contradicts the ideal case described earlier, where would prefer:
d 0 dy2..0 = -dg..0 -dgn..
Vqe (0, k + +1, ,l, Vie (k................ nt),
(6)
and d 20
51

---

Towards this, we first define a row-summed quantity:
dg gmml e
d
gm
4 m
To satisfy the condition in the ideal case (Eq. 6), we now introduce a preconditioner based optimizer:
P++1 P 7t D; V V.. L(6)
where
Pr
[Po,
Pı.t
pn,r]'
BR++1
dorat dgid
doi.t Age
D diag
0, 1, 1 Aga.t
1....,, e R (x+ 1(++1)x(++1)
Here, Pi,t denotes the ith positional encoding scalar at iteration t. Therefore, after the first update step (i.e.,
0), we obtain:
P1,1 P2.1 Pa, P4. Vg (0, ......n
Note that while we assumed = -oC for analysis, in practice we can set it to a large negative constant to ensure that -dgm. >> We need to show that this property holds for all iterates. By Lemma 3, we can show that for each iterate of the preconditioned descent, the following holds:
Pl.t P2.1
Pk.t
Pq. Vg E {0, + 1, ...
Furthermore, by Lemma and Lemma 4, we observe that P1 and p2 tend to infinity as they increase monotonically over time, whereas the other positional scalars either tend to -0 converge to constant. addition, Lemma 4 provides the approximate time required for the iterates to reach the optimal attention configuration. Therefore, by the end of the first stage of training (as T₁ o0), can recover the ideal attention map corresponding to 1tth -order Markov chain:
Vi ((e, + Cx, Ex.-E)
By reintroducing the layer normalization and taking the limit CC, recover the conditional k-gram estimator, as shown in Appendix A.
52

---

E Training Dynamics: Stage One Lemmas
We the following notation:
logit [X, softmax (azZn)]
2 ez
eg' X, softmax softmax (1(1) X A (2) 1
Σ exn--
(8)
Initially, we know that A 0220 I Hence, we can rewrite this as:
logit. sc. Xn softmax 02,0 softmax (A(1)
X.
Σ Ezn
(9)
Te
(X.:s)
Lemma 2 (Exension of Lemma G.I. in Nichani al. [20] to hthh -oddrr. Let (2,1)
a2,0!)
0 (A for 02,0 & Define g 9s E R$ by
Xn ((48 (Xnik))e X DD exe cmsi13
gt : n Σ P(-)(())
$1,'.., 8k)
logit, .n.0 +E
s', 1118
e D(A (Xn; fnie))); X7 (EN=O
k (ean-15544)
L
91
7(8' 11,11111.11,
logit, 164 m. +6
Then, holds that
1 Dile & ( 1) & 6
Proof. The proof follows similarly to the proof of Lemma G.1 in Nichani et al. [20]. We begin by bounding the difference inside the expectation.
We first apply the triangle inequality separate the difference between the two fractions:
e TT,,D,DA, (Xn; ((x)
((A (Xn;
logit " n. +c
logit n. +
(eA (Xn:k))e
logit bp 0. +c
logit
n n +e
|eJ.Kn (X(:)))(xnn1))) - ed
logit s, .n.0 +e
To bound this expression, we begin by bounding the logit difference in the numerator of the first term
logit s nn. 0 logit
n
e (A, (2) (Knik) - A (Xni k))
< l Xn la 1446'(Xnnk 1- (Xn;
A
(2) (Xni k) - 1/8) (Xni k))ll1.
53

---

where the last inequality uses the fact that |ellxx & 1. We now recall the definition of Ag) and analyze its properties.
k-1
A (Xn:k)) = softmax 22.0 (alimax((44))a
Σ
Cxa-t
In particular, if 02.0 0, the argument to softmax zero, and we obtain the uniform distribution:
A (Xn; softmax(0)
Since the entries of the softmax input are in [0, 1]. we can bound the minimum and maximum values of any softmax output entry as follows:
2
1) + & A (Xn; E
( 1) + 602,0
These bounds let now quantify the maximum possible deviation between corresponding softmax outputs.
sup Al (Xn; k)); A kX):
22.0
(n 1) + 2
71
(n 1) 0200 1)
n((n 1) + c02,0)
6 1
71
Similarly, for the lower bound:
inn (X -A, (x k));
1
(n
1) 1)20 +
n
220 1
Thus, we conclude:
(xx; - A (X.:
Next, we simplify the derivative expression using the identity for the softmax Jacobian.
(X2 (a (x,; - logit,, e,, .n.e)
We now bound the difference in these derivative expressions using triangle inequality:
le,x, (X,: k))e - (X,, k))es < (4 (X,: k)₂ - 1 (X,: k)) [l,, - logit m.o m.
54

---

Al (2) ())
|logit 5, n. logit 5/ _,.
220 -1
< 2.
Here we used that - logit & and Ao < 1.
Combining the bounds for the numerator and denominator differences, we obtain:
X, (Xm:)((x;)))))))
(Xn;k))es
e(A (Xn: k:)le
202.0
logit % n0, +C
logit n te
nn
Finally, we bound the overall error 1
golloo by summing over all possible state sequences:
3 1)
11,, Delloo & n
Σ
TTss 81, s4)
2n
"'," 18k
(c - 1)
= 3
Σ
Eg (T(s' s1. ,kh
22
$ 35 .2,
(E -1 1)
Using the inequality ea2. 1 & 202,0 for 02,0 E [0. 11, we conclude:
l1 - Dillo & 6
Lemma 3 (Extension of Lemma D.3 in Nichani et al. [20]). Let A 02,01. There exist constants -3/- Cy Sn such that, 0220 & c,,((n(1 A)) 3/2 Ifr I, where I 11, k}, where k denotes the order of the markov chain:
Cr,sa2.
G (A (A(( (0)))) (ofttax(AA
softmax (A())) (1
n
Note that the subscript in this case denotes the th row and the subscript i. denotes the j"h element in the tt row.
Proof. The proof follows similarly to Nichani et [20]. The gradient (())(A (()))))(()2))))))))))))))))))( A(2)); can be expanded as
(A() -a2.oD(softmaxt (A()))
Σ
Ba,x [T(S' s1, ,8) an (X)].
where
Es (x logit te e ((A (Xm:)
(Xn;
Σ k (to--1
1
Define
0,¹ :
Σ
Ex-P(r) [r(s' s1,".. sk) Eor sx (X)].
g. == n Σ ExaPix) 7(5' s1. ,kx) Eo (x)
2',1 88
55

---

where (X) is evaluated at (A(())).
Thus,
(A()))))))) = 02.0 ((softmax(AA.))
Sn
Since 020 & c,,s(n(1 - 1))-3/2 and applying Lemma 2, we obtain
-3/2
C
16 - 1
(n(1-X)
For some arbitrary constant C. Thus, it suffices to analyze Ji.
Computation of 9: under
Under 0
A (Xnik) - In
thus,
logit = fsx(('))
Therefore,
e e (x); = -In-ss--aa (xx(s)))
Thus, the j-th coordinate satisfies
T 81, ,8k) (Im-ss (xx(s))) 81+1
Σ E P(m)
(px((5)+ +
Applying Lemma
C,.8
19., -isl &
(n(1-1)
By Lemma 7 (for first order Markov chains) or by (App. D.2.1, Lemma. 5 and Assump. for second-order Markov chains) by (App. D.2.2, Lemma. 6 and Assump. 3 for any-order Markov chains), for all j # T I,
gh. Dir & $
9ir
1gi.
Jigl + (gt,p())
-
S
$ 5
Vn(1 Aj
< 8/2.
Note that, in the last inequality, assume be long enough for this hold.
Now expanding:
02,0
(sottmax (ottmax(A(())). (A(2)) - (softmax(A (A g softmax( (A
Sn
(4)())) ((2)
02,0
Sn Softmax(A!)) ) - (softmax(A (T) g, softmax( A (()))))
Thus the difference is
(A() -19))))))))((()))) (22)))
56

---

02,0 (softmax(44) (A coftmax(A!")) (0
(softmax(A (A 9.))
+ (5 - -hhsss
Vre
Applying bounds:
gi Pi.j 2'
thus,
42.0 (otmax(A!"))dl (atmaa(44))) 8 Sn
)
Lemma 4 (Extension of Lemma D.5 in [20]: Bounding the time of convergence). Let (
a2 where 02 , 3/2 >) 3/2, Ler {1. .k.. Then there exists
( - (n (n-k))
mlog(2n) + a loo
1 1 k ( (t)
such that for any x TI
softmax(A > --a
for all i and all I.
Proof. The proof follows similarly to the proof of Lemma D.5 in Nichani et al. [20]. Let I = Then, by Lemma 3, we have that throughout training,
(L....,k). k).
(1), 2 (1)ss for all j $ I,
for all I.
Moreover, by Lemma 3, the quantity softmax (t)) monotonically increasing in for all r e I.
For fixed i, define the gap function:
A(t) (t),r max (1) JØI
By properties of the softmax function, we can compute the lower bound:
exp(A(())
(tt)ma(((((t)ar> n kcxp(A(t)
Now, consider the first time T+ (1/2) such that
sotmaa((a")(((112))))+++) 1
2k
For all < (1/2), we have
1
softmax(A()()))
21:
Thus, by Lemma 3, the gap A(t) evolves
C,,502.0
A + 1) A(t)
n²
57

---

Consequently, by successive addition over time,
02.07
((T2(1 Z
T+ (1/2).
n²
Suppose for contradiction that
A(T+(1/2)) >> log((n).
Then we have
exp(log(2n) (2n))
softmax (A (T+())))) 3
11 + k eeep 2n))
2n
+ 2kn
2
2k + k/n
Since > and n is large, we have
2k 1
$ 4k,
and thus
>
2k + 1 k/n
2k
This contradicts the definition of T+ (1/2) as the first time when softmax (1))r must have
softmax(A1 >
1. Therefore,
2x-
A ( (1/2)) < log((n),
and hence
(1/2) & nnnnjiaaa, 22.0 log(2n).
Now define T+ (a) as the first time / such that
()ftmex(4(( 1 a, fol e I.
This implies that, for sufficiently small positive a, to reach contradiction, would require:
(a) n 1: + k- (a)+
a,
! a) (n k) a)
((T+(a)) > log
e (4+()) the gap (t) evolves according to
A [ 3 A(t) C4,sa2,0a
Therefore, in order to achieve the required increase in A(t), we must satisfy
( a) (n k)
T+ - ((a))-+(1//) & na na- -a-t log
1 1 a))
Combining with the earlier phase where
7+ (1/2) $ n²n -a,0 log(2n),
58

---

we conclude
(t -a) (n k) (#= (2)
T4(a) & n'nn" 020
2 log(2n) na -nantnn-- ' log
Grouping terms, we obtain
( a) (n k:)
T+(a)= & uyni a n nlog(2n) + aa--oog
1 1 k ( a))
59

---

Higher Order Markov Chain Lemmas
Lemma (Successive Markov Chain Lemmas), For a second-order Markov Chain, = | 0) T( 1, e 8' and -[1[1,1) - π () >> then the probability P (0 0, 0, and Pa (1 1,1)
decreases with increasing Moreover,
Ex-P (P
En-P [P* 0,00 (* 1,,1)
Ak > 4
Proof. We begin by establishing the condition for two steps and then generalize by induction.
P² 0, 0) = x(0 0, (0 0,0) + =(0 1, 7) | 0000
Given the assumption r(0 00,0 > m(0 1, 0). we multiply both sides by 1 x 0,00):
<0 0,0) (1 (0 0,0) > >(0
=(0 0,0) > 7(0 0,0) + 1, | 0,
(0 0,0)
(by definition)
Now, we apply induction.
(0 0,0) =(0
0.0) 000000000000000000000011100 ( [10
(0 1,0) + (,(0 -(0 ( 1,0)) (,00
< (0 1,0) + (z(0
+.(0)(=(000)))) 0,0) - -(0 1,0)) x(0 0,00,,0)
0,000
Hence, by induction:
pi+1(0 0,00 = x(0 1,0) + (=(000,0) #(0 1,0)) P'(0) #0 1,0) (=(0) 0, m(0 1,0) p²-1(010.0)
< (0 0,0)
Now, using the recursive relation derived above, we express (0 0, explicitly. The recurrence is:
pi+ 00
0,00 T 1,0) + (m(0 0, =(0 1,0) P5(0 0,0)
We separate this into a homogeneous and a particular solution:
Homogeneous Solution:
P (0 (T 0,0) =(0 1,0)) 0,
PA 0,0) = C' (=(0) 0.0) m(0 1,0))*
Particular Solution:
P. 700 1,0) (=(000,00 #(0 1,0)) P,
60

---

<<0 1.0)
P3 m 0,00 = 1,0)
Using the initial condition P0 0, 0) #T(0 0, 0), we solve for the constant:
(0 0. 0) P.
C= (0 0,0) - x 1,0)
Final Solution
T(0 1,0)
PC(0 0,0) = 1- #(0 0,00 =(0 1,0)
T(0 1,0)
+ T 0,0) 1- (0 0,0) + (00 1,0)
(=(0)0,0) =(0 1,0))*
Similarly, we obtain:
pitl (((111))) VieN
Now, we can first express
T1 0,
( 1- 1,1) m
| 0,
7(1 - (1 1, (11 0,11)
(=(1]1,1) m(1 0, 1))
1-
Now, define
x(0 1,0)
G. T(0 0,0 1- [- 0, + 7(0 1,0)
| 0.
G
(T( 1,1) - 1- (1 1, (1)
Thus, we can write
TT0 1,00
(9 0,00
+ Go (T(O 0.00 a(0 1,,))²
1- (0 0,0) 7(0 1,0)
T(0 1,0)
p0(0| 0,0) 1- =(0 0,0) 700 1,0)
+ Go (T(( 0,00 =(0 1,0))²
and therefore
pd 0, (0 0,0) G. (T(( 0,00 x(0 (1 (1- (T(( 0.0) =(0 1-2)
We further simplify Go:
=(0 1,00
Go (0 0,00 1- (0 0,0) +00 1,0)
(1-x() =(0 0.0) + 7(0 (=(0
0. ((),,))))(((
T(0 1,0))
=00
- =00
1-
61

---

By assumption, x,0) =(0 1,0) > 8' and 1 x00 0,0) > é', thus
Go e"
for some positive constant
Hence, substituting back, we obtain
(O 0,0) (1(0 0,0) 8m
(Where 8 is positive constant)
Taking expectation over π we have
E (0] 0, - (1( 0,0) > 6,
where 8 is defined formally as follows:
The constant 5 represents a uniform lower bound on the expected gap between the third-step and Vi-th-step transition probabilities across all Markov chains sampled from the prior distribution P. is defined as
8 inf E (0 0.0) P0( 0, 0,0)].
x~P
and satisfies 8 > This ensures that, even after taking expectations, the separation between transition probabilities remains bounded away from zero across all sampled chains.
Similarly, for the other case, obtain
PS 1 - P'(1 1,1) > 8"".
E [1 1,11 - P'(1 1, >
Therefore, using linearity of expectation, we conclude
Ex 0,0) + P3 1, E, P (.00) 0.0) + P (1]1,1) 28 A4 4.
Lemma (Monotonicity via spectral decomposition). Let (Xt) reversible, irreducible, aperiodic kih-order Markov chain stale-space S, and let Sk. For integers i, 2 define the return count
Σ
Σ
Pr(Xite= = 8j | X =8).
9i.
(SL....N))SS
1=1
Equivalently, one can view the chain as a lifted first-order markov chain on A, where a = (a),.... ak), (b.....,bu) and
P(a b) Pr(((.,..... (1-xx+1) b (X.-1...., a).
f(a) |a| (where |. denotes the cardinality),
Ii.E = Σ b)
a.DEA
Under detailed balance with stationary distribution т. If P self-adjoint on e πt) with eigenvalues X1 X2 2 ... AN > and if Bm > 0 where orthonormal eigenfunctions are defined as {6m) and 8m is defined as:
Bm Σ abEd
(ml()) 7(6) f(a).
Then for all t 0:
N Σ
2 (1 )mm 2 0,
- 9i.41
m=2
62

---

N
Σ
(2
2mmmm 0.
91,1 Di
Hence, 9t,1 > 91,2 > 91,3
Moreover, 91.1 3 2h-1 if ΣN - 2m 3 2
Proof. From spectral expansion, we know that:
Pa(a ((a-b) b) Σ X pm (m)) (b)
Therefore, we obtain (on substitution)
Σ
1at Bm
9i.E
=
Hence, finally obtain monotonicity, due to the assumptions:
Σ
x1+1(1 Am) 2 0
gi. gi
m32
The last condition obtained by direct substitution.
63

---

G
Additional Lemmas
Lemma 7 (Lemma D.2 in Nichani et al. [20]). For all j 3 0. j t 1, we have
9t,1 gi 7
25
Where:
(i)(T) Σ ITs' 5)
P(xx = ,
Fi = 8) - -1,
(m((s))
9i.1 (T))((r) [gi,(
Lemma 8 (Levin and Peres [13]). Let (z)) be stationary, reversible, ergodic Markov chain with stationary distribution n. and let denote its Markov operator (T). For any function f: Z - R with ExIf*] < ap and all t,t' 3 O,
Cor(f(z), (J, (J,plt-r1
where J = - and h), == (=(o(>)/(=))-
Proof. By stationarity and the Markov property,
E[f(ze) 21 = (p' 1pt--tpp -f)(z))
(10)
[=[[(=)[(E+)] = (J, p'
(11)
Thus,
(fo(f())) = (f, pl A)x (En(f))?
(12)
Expanding f F F. and using that preserves constants,
(1. plt-
D> P
(l, + (E+)
(13)
which gives
(o)) (F, F).
(14)
Applying Cauchy-Schwarz Schwarz inequality,
((or(f()), & | [(()))
(15)
If the spectral gap satisfies |Ploop & Å on centered functions, then
plr- M22a()) & x1e
(16)
thus
(Co)((()) & var Va)=(f).
(17)
Lemma (Variance Lemma similar to Cor. F.6 Nichani et al. [20]). Suppose:
1. (st)aa is stationary, irreducible, aperiodic k-th order Markov chain on finite state space S size S,
2. Define the lifted process 2, (St, St-1 St kp S*.
3. Assume the transition matrix P of (zt) has second-largest eigenvalue modulus À 1.
Then,
px(('))
I,=='
n2 1-1
is the empirical frequency of observing 8', we have the variance bound
E
[jx((s')
(n())))]
1
n(1 x)
where:
64

---

(x(s') (=(3))=T((2222) s'}) is the stationary probability of s',
and the hidden constant is universal (independent of
Proof. Let 2n (In, "n-1."""""""
Therefore, we can see that:
(n, = Σ 1
(18)
Therefore:
Var(pn.x(s)) = Var
(19)
Cov (
(20)
From Lemma 8, we can see that,
Var((a,.xss))
Σ x1--1
(21)
π1t
(22)
n 1 Xj
Lemma 10 (Extension-1 to -order of Lemma G.2 in Nichani et al. [20]). For any 8, S and
any tthh -odder transition kernel, æ with spectral gap X(T) X, and with pz(s)) 2 0, and
i >> jthere exists sufficiently large constant Cy,s such that ife Cr.5 (n(1 A)) 1/2
A)) then:
(I Px (x((s')) (xx ix(8)) +6
= 8, (Px]
Px(s; 01)
Ex
(x((8)
(n(1-))
Proof. The proof the same as the proof of Lemma G.2 Nichani al. [20]. where control the variance term, we use 9, For completeness, we desribe the proof below:
Define
(Tr,-s' (x)) Ax(s') +
Px[x, 8, s)
(s(ss)) = Ex
FX s).
(x()))
Expanding:
[Ir,= - - In.-s - pxx (s')Iz,=x
Thus:
be,- Px[x; = 810,- s]
ux (x)s'))le,== (s') +e
(s, = Ex
Ex
+ Pxx1z = s].
fx (s' +e
pxx (s")
Grouping terms:
(s, Ex T -ly (xx(s)) (px(S))
1 m1167) )
+ Fx Tzy- (1
Ax(s')) + )
xx((')
Grouping terms:
(=(8)) (5) (66))-6 -
(s') +e u=(s")
(5x + (s')
65

---

and
px(s") Ax(s') te ix(s')
Thus by the triangle inequality and boundedness of indicator variables Tz, Tzy-s E {0, 1J, we get
I N===rl;==llāx((( (x((81) + e e(Il,, + (
(s, x Ex
(s' + c))
$ Ex [ix((')) Mn((')) +e P (xx(6)) (s))
Mx(s')
2
Using concentration inequalities (Lemma 9):
Px ix((s)) & Mr(s')
Ex [(xx((') - (=m(s')) n(1 -) X)
x
2
n(1 - X)
and (n(1 2))-1/2 we conclude:
1 (x, $ n(1 X)
Lemma 11 (Extension-22 -order of Lemma G.2 Nichani al. [20]). For any s' s', 82. 8k
and any Kth order transition kernel, with spectral gap - X(T) > - X, and with ( > 0,
and j. there exists sufficiently large constant C,,5 such that if C. (n(1 A)) 1))-1/2 then: 1/2
(Lo:-o (
]- (Fxls; -s,
Px|s; = sil) $ vn(1- A)
Ex
ix(s') +E
(s
Proof. The proof follows immediatly from Lemma 10. Let:
(Tn=-' ( ix((d)) +e
[x[ss 5' st]
Px|s; = sil) = vn(1 A)
Ex
Mr (s')
Qi
We know from Lemma that Qi $ Vn(1-x) and hence immediatly follows that Σt 1 x
Vi( X)
66

---

H Experiments and Experimental Details
H.1 Additional Experiments
In this section, we present additional attention map visualizations obtained by training a two-layer, single-head transformer on first, second, and third-order Markov chains, respectively, Attention maps from both the first and second layers are shown. For the first layer, we compute the average attention map over multiple sequences to understand the expected attention pattern. For the second layer, we display the attention map for a specific sequence. Additionally, we include Pseudo Attention Map, which highlights the tokens that would ideally be attended to under k k-gram estimator, To assess alignment, we also plot the absolute difference between the actual attention map and the pseudo attention map. Note that for the first layer, we do not compute pseudo attention map; hence, pseudo attention and the corresponding absolute difference are not meaningful in that context.
H.1.1 First-Order Markov Chains
The following plots show the attention maps learned from first-order Markov chains.
70110
Figure 5: Layer 1: Average attention map computed over multiple sequences.
Misrica patorrci Seige Rusdi stastos
Figure 6: Layer 2: Attention map corresponding to a randomly sampled sequence (denoted as sequence-I)
Figure 7: Layer 2: Attention map corresponding to randomly sampled sequence (denoted as sequence-2)
67

---

Figure Layer Attention map corresponding to randomly sampled sequence (denoted as sequence-3)
H.1.2 Second-Order Markov Chains
The following plots show the attention maps learned from second-ordes Markov chains.
Figure 9: Layer I: Average attention map computed over multiple sequences.
Figure 10: Layer 2: Attention map corresponding to randomly sampled sequence (denoted as sequence-l)
Figure 11: Layer 2: Attention map corresponding randomly sampled sequence (denoted as sequence-2)
68

---

Figure 12: Layer Attention map corresponding randomly sampled sequence (denoted as sequence-3)
H.1.3 Third-Order Markov Chains
The following plots show the attention maps learned from third-order Markov chains.
Figure 13: Layer 1: Average attention map computed over multiple sequences.
Figure 14: Layer 2: Attention map corresponding to randomly sampled sequence (denoted as sequence-l)
Figure 15: Layer 2: Attention map corresponding randomly sampled sequence (denoted as sequence-2)
69

---

Figure 16: Layer Attention map corresponding randomly sampled sequence (denoted as sequence-3)
H.2 Experiments on One Layer Transformers
To assess whether single-layer transformers can implement induction heads, trained one-layer models on sequences generated from second-order Markov chains, using the same hyperparameter ranges as in two-layer baselines.
Quantitative results. As shown in Tab. 2. the gap between the true loss (computed from the ground-truth transition probabilities) and the loss achieved by the trained model remains significantly larger for one-layer transformers than for their two-layer counterparts. This difference persists across hyperparameter settings and indicates that deeper models capture the underlying structure more effectively.
Number of Layers
Excess Cross-Entropy Loss
0.131 0.000
0.100 + 0.003
Table 2: Excess cross-entropy loss (loss obtained by the model subtracted with the true loss) for one and two-layer transformers trained on sequences from second-order Markov chain.
Qualitative results. Examination of the attention maps (Fig. 17, Fig. 18, Fig. 19) confirms this finding: one-layer transformers converge a qualitatively different structure than the pseudo-attention map required for a k-gram estimator (see Fig. 1b). Specifically, attention mass concentrates along the two three lower diagonals, rather than following the induction-like pattern.
These results are consistent with prior observations [22] that one-layer transformers are insufficient for representing induction heads, while two-layer models succeed in doing (albeit different datasets).
Figure 17: Attention map corresponding to a randomly sampled sequence (denoted as sequence-l)
70

---

Figure 18: Attention map corresponding to randomly sampled sequence (denoted as sequence
Figure 19; Attention map corresponding to randomly sampled sequence (denoted as sequence-3)
Figure 20: Attention map corresponding to randomly sampled sequence (denoted as sequence-44
H.3 Experiments on Noisy Sequences
We evaluate the robustness of our k-gram framework under non-ideal conditions by introducing stochastic noise into synthetic sequences generated from recond-order Markov chains. We consider two noise processes: (i) random token substitution and (ii) perturbation of the transition dynamics. Models are two-layer, single-head transformers trained with three random seeds. Performance is reported the difference between the model's cross-entropy loss and the cross-entropy computed under the true transition probabilities over a large evaluation batch.
H.3.1 Random Token Substitution
We corrupt sequences at the token level with probability Pn- For each token, with probability pa, we replace it with a uniformly sampled alternative; with probability 1 Pa it is left unchanged. The results can be found in Tab. 3. In the noiseless case (p, = 0), the learned second-layer attention map closely aligns with the greudo-attentinn map (the optimal attention pattern for ideal k-gram estimator). As Pa. increases, the alignment degrades and the loss gap widens.
H.3.2 Perturbation of Transition Dynamics
We next perturb the sequence dynamics directly. Let Pa denote the true transition probability matrix. Before each prediction step we interpolate between Pa and randomly generated transition matrix P.. defining
P (1-a) + aPr. a [0.1]. 1],
(23)
71

---

Pn
Difference in Cross-Entropy Loss
0.0
0.10 + 0.003
0.1
0.15 + 0.010
0.25
0.18 + 0.002
0.5
0.21 + 0.005
Table 3: Cross-entropy loss difference under token substitution noise. Values are mean = std. across
three seeds.
Ck
Difference Cross-Entropy Loss
0.0
0.100 + 0.003
0.1
0.101 + 0.003
0,25
0.141 + 0.005
0.5
0.181 + 0.004
Table 4: Cross-entropy loss difference under transition perturbatioon defined in Eq. (23). Values mean + std. across three seeds.
where a new transition matrix P. is resampled at every time step. As in the previous setting, we measure the difference between the model's cross-entropy and the true loss. Note that the results can be found in Tab. 4. The loss gap increases with a. For a & 0.1, attention maps remain close to the secudoattention pattern; at a 0.25 and beyond, alignment deteriorates substantially. The model is more robust to this transition noise than to token-level corruption.
H.3.3 Discussion
These results complement our theory: when the iata-generating process is close Markovian, the learned attention structure approximates the pseudo-attention map; as noise increases, alignment degrades tandem with cross-entropy performance. Overall, transformer-base estimators are sensitive to deviations from ideal k-gram dependencies, with greater robustness observed for transition perturbations than for token substitution.
H.4 Experiments on Bit Precision
Our bit-precision results are the same as those of [26]. which show that [2olog + k) bits per parameter suffice for an additive error of O(1TT), where T' is the sequence length and l the Markov order. To validate this, we quantize weights and activations of a model trained on second-orde Markov chain with sequence length 32. We then compute the Frobenius norm between the pseudo (optimal) attention map and the second-layer attention map of the quantized model across 10 random sequences (each from different Markov transition kernel), and report the mean and variance:
Quantization (Bits)
Frobenius Norm (Mean t Std)
2
4.6108 - 0.7190
4.5842 = 0.9826
8
3.8615 + 0.4590
(No quantization)
3.7700 + 0.44255
These results suggest that 8-bit quantization preserves the attention structure (as measured by the Frobenius norm) almost as well as full precision (32-bit), and aligns with the theoretical bounds.
H.S Experimental Details
The table below provides details of the hyperparameters used across all experiments in the paper. Fach experiment was conducted on single NVIDIA A A100 GPU, with runtimes ranging from 30 minutes to hours.
72

---

Dataset Architecture
k-th order binary Markov souree
Based the GPT-2 architecture implemented [24]
Batch size
Grid-searched in [16,32,61] 641
Accumulation steps
Optimizer
AdamW (81 {0.85, 0.9), 82 E (0.9,0.95). Grid-searched {0.0001, 0.001, Cosine
Learning rate
Scheduler
# Iterations
30000
Weight decay
Grid-searched {0, 10 1001000..00000 (010.00-))
Dropout
or
Sequence length
32, 64
Embedding dimension
Grid-scarched in [32, 64, 128)
Transformer layers
2
Attention heads
1 or depending the experiment
Repetitions
5
Table 5: Settings and parameters for the transformer model used in the experiments.
73

---

I Parameter Breakdown and Comparison
In this section, we present comprehensive comparison of parameter counts that explicitly accounts for all components of the models under consideration. Our analysis shows that our construction achieves greater parameter efficiency than the design of Rajaraman et al. [26] and also improves upon the approach proposed by Nichani et al. [20]. To make this concrete, we provide full parameter breakdown for both our model and that of Rajaraman et al. [26], detailing the contributions of embeddings, attention layers, feed-forward modules, and output projection.
Our Construction
The total parameter count of our construction is: 9 + 3)2 (68 (2T + 25 + 9)
Component
Parameters
Token Embedding
(68 + 3) S
Attention Layer
Positional Encoding
(65 T 366 +32
Query, Key, Value Projections
MI.P (Three Layers)
Weights
36S 3)2
Biases
3 + 3)
Layer Norms
6 + 3)
Attention Layer
Positional Encoding Query. Key, Value Projections
(6S + 3) T
3(65 3]²
Output Projection
(68 + 3) S
Total
9(68 + 3) + (6.5 + 32T + 28 9)
Table 6: Parameter breakdown for our construction.
Construction in Rajaraman et al. [26]
The total parameter count of the construction in Rajaraman et al. [26] is: 15(6S + 3)² + (65 + (3T + 2S + 10)
Component
Parameters
Token Embedding
(6S + 3) S
Transformer Layer
Positional Eneoding
(65 + 3) T
Query, Key, Value Projections
3(6S +3]2
MLP (2 layers) Weights
266S + 3)²
MLP (2 layers) Biases
2 + 3)
Layer Norms
2(6S + 3)
Transformer Layer
Same as Transformer Layer 1: (6S + 3) x + 5(65 + 3)² + 4(6S + 3)
Transformer Layer
Positional Encoding
(65 + 3) T
Query, Key, Value Projections MLP (2 layers) Weights
3(6 +3
266S +32
MLP (2 layers) Biases
2 + 3)
Layer Norms
2(65 + 3)
Output Projection
(6S + 3) x
Total
15(65 + 3) + (65 + 3)(3 + 28 + 10)
Table 7: Parameter breakdown for the construction in Rajaraman et al. [26].
74

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
All claims are thoroughly addressed in the main text.
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 be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations and future work are discussed in the Introduction and Conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those 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 well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach 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 in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
75

---

Complete proofs of the theorems of Sections 3, 4, and 5 are provided in the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All implementation details and hyperparameters are provided in the Appendix. A link to the complete code used is provided in the abstract.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is 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, it 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 can also be provided via detailed 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 appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a 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 be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
76

---

Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: A link to the complete code and instructions run is included in the abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
ccips.cc/pbblcc/eeee heppliii/ggide../.ee /CodeSubmi ssi agionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: All training details and hyperparameters are included in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can 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: All empirical results have been obtained via multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
77

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but 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 of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The experiments of the paper require minimal computation and memory capacity. Details have been included in the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: pettpa:/needee /ttp::///eeu..eeeeeeee httpppic/ttiissuiiiiiiieaaao
Answer: [Yes]
Justification: Ours is mainly theoretical analysis, and the NeurIPS Code of Ethics has been fully respected.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consideration due laws or regulations in their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification; Ours is mainly theoretical work which should not have any important societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
78

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 partieular 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 point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
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)?
Answer: [NA]
Justification: Our model is trained from scratch from synthetic data, which does not pose any risk.
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 usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The original authors of the considered model are properly mentioned and cited in the main text.
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 URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
79

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode wwwccoo///////// datasets 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 it 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 are introduced in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. 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:
The paper does not involve crowdsearching research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing 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.
80

---

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.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development this research did not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy htaps::/nneuris... http::pppp.p.pooeeeeee////////// for what should or should not be described.
81


Reviews:
[{"review_id": "Fgmz0TVezn", "reviewer": "Reviewer_2Q6Z", "summary": "The paper shows how a k-th order induction head can be represented by a 2-layer transformer with a single attention head and MLP in each layer. The key idea is that the tokens can be expanded to a sufficiently high dimensional space where they are sparse (one-hot encoded, in this case) and attention weights can be used to combine them in superposition with an exponential moving average -- to represent k-grams. Prior constructions required two different moving averages (of window length k) shifted by one, but this paper shows how one can be used to derive the other since the token embedding at a timestep $x_{n-k}$ can be estimated using a suitable bias & ReLU on the moving average quantity $v_n$, and these can then be linearly combined to estimate $u_n$. In the second layer, the similarity between $u_T$ and $v_n$ is used to trigger an induction-based completion mechanism for $x_{T+1}$ copying $x_{n}$. While that shows that k-th order Markov chains are representable, it leaves open the question of whether they can actually be discovered during training by optimizing the cross-entropy loss on tokens. For this, the paper performs a perturbative analysis of stability of the proposed architecture weights (actually a subset, thereof) to show how optimization (at least within this subspace) converges to the anticipated weight patterns. The paper also points out an apparent \"depth-width tradeoff\" given that a k-gram model could be represented with a 3-layer transformer with a single attention head in each layer or alternatively a 2-layer transformer with two attention heads in the first layer.", "questions": "1. Why should the attention weights have the form $3^i$ instead of any other scale for the exponential? Is this precise number validated by experiment, or only the exponential behavior?\n2. In Figure 2, layer 1 -- might it be interesting to have an additional plot depicting the attention map intensities on a logarithmic scale? Or maybe plot the variation across different diagnoals and inspect whether the weights are exponential as proposed (or whether there might be some other interesting pattern).\n3. In Appendix E (fig 5/9/13), why not compute a pseudo-attention map for the first layer?\n4. In Fig 2 Layer 2 (likewise figs 6-15) -- is it correct to say that the first \"k\" columns should actually be ignored when comparing with the pseudo-attention because the k-th order context is actually ill-defined for the first k tokens? If so, inspecting the \"difference\" between the \"learned attention\" and \"Optimal\"/\"pseudo\" attention might be misleading. The comparison might be less misleading (and also more favorable!) if you consider a sequence of length k+N and showed the attention matrix for a window of length N starting with the (k+1)-th token. Would you agree that's a more interesting plot?\n5. Have you considered the implications for inference with quantized weights/activations? This proposal might point to specific patterns of failure, which could then be compared against LLMs for more circumstantial evidence on the representations within LLMs.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n1. The proposal is conceptually elegant, including the minimality in components needed to represent k-th order Markov chains.\n2. The experimental results seem closely matched with the theoretical model.\n3. Explores not just in-principle representability, but also feasibility of discovery by training/optimization.\n4. Such principled understanding will hopefully have a positive impact on architectures and training methods over time.\n\n## Weaknesses\n1. With this framework, the largest allowed \"k\" will be practically limited by the finite-precision of activations. This is even more of a problem when inference uses quantized activations. [addressed in rebuttal]\n2. The observation regarding the depth-width tradeoff seems quite weak / superficial, and of unclear significance. It would be more interesting if the authors could flesh this out further. [partially addressed in rebuttal]", "was_revised": true, "final_justification": "I am satisfied with the clarifications provided by the authors in the rebuttal. I believe that adding all those details to the final version of the paper would strengthen in -- making it more interesting and also more readable. On further reflection, I am happy to raise the score (from 4 to 5) for the promised improved/final version.", "rebuttal": "We sincerely appreciate the reviewer’s thoughtful comments, valuable suggestions, and recognition of our work. Below, we provide detailed responses to each of the individual comments and questions.\n\n---\n\n## Comments about Weaknesses\n\n**W1. With this framework, the largest allowed \"k\" will be practically limited by this …  quantized activations.**\n\nWe agree with the reviewers perspective. To clarify and provide additional context, the bit precision in our representation results follows prior work on constructions for the Markov in-context learning setting, namely Rajaraman et al. [1], where $\\Omega(\\log T + k)$ bits per parameter are sufficient to achieve an additive approximation error of $O(1/T)$. Where $T$ denotes the sequence length and $k$ denotes the order of the markov chain. \n\n**W2. The observation regarding the depth-width tradeoff seems quite weak / superficial, and of unclear significance. It would be more interesting if the authors could flesh this out further.**\n\nWe agree that the most compelling result is the two-layer, single-head construction. The depth-width tradeoff presented in Section 3 serves two purposes: first, to provide context and continuity with prior work, thereby setting the stage for the subsequent constructive proofs; and second, to highlight a potentially interesting direction for the community. Specifically, the ability to trade depth for width may offer opportunities for inference-time speedups in other settings. While the two-head, two-layer construction may appear redundant for the specific task considered, we believe that exploring such architectural variations could prove valuable in broader applications.\n\n--- \n\n## Questions\n\n**Q1. & Q2.   Why should the attention weights have the form ... only the exponential behavior? In Figure 2, layer 1 -- might it be interesting  ... whether the weights are exponential as proposed?**\n\nWe address both the questions together below. We would like to note that there may exist multiple possible constructions demonstrating how two-layer, single-head transformers can represent induction heads due to the highly non-convex nature of the problem; our approach is one such construction. Our attention map design is motivated by two key observations:\n1. *Empirical Observation*: In the first layer, we consistently observe that the attention weights increase monotonically along the first *k* lower diagonals and sharply decay to near-zero beyond that (figure - 2 in the paper).\n2. *Prior Theoretical Construction*: This stems from a theoretical construction described in [1], where the attention pattern constructed in that paper is proportional to $2^i$ (a dyadic sum representation).\n\nFollowing the reviewer’s valuable suggestion to rigorously examine whether the diagonals in the first layer actually adhere to this pattern, we conducted further experiments. Specifically, we trained two models: one on a second-order Markov chain and another on a third-order Markov chain. For each model, we sampled 10 sequences and extracted the lower-*k* diagonals ($k=2$ or $k=3$) of the first-layer attention maps (note that $k$ refers to the order of the markov chain). For each diagonal and across the sampled sequences, we computed both the mean and the standard deviation of the attention values. Below, we summarize the results for clarity. *(Note: L-k refers to the k-th diagonal below the main diagonal.)*\n\n### Table -1\n\n| Order | L-2 Mean ± StdDev| L-1 Mean ± StdDev| Ratio L-2 / L-1 |\n|-|-|-|-|\n| 2| 0.514 ± 0.0006| 0.256 ± 6.2455e-05| 2.00394|\n\n| Order | L-3 Mean ± StdDev| L-2 Mean ± StdDev| L-1 Mean ± StdDev| Ratio L-2/L-1 | Ratio L-3/L-1 |\n|-|-|-|-|-|-|\n| 3| 0.510 ± 0.0033| 0.3012 ± 0.00068039 |  0.1122 ± 0.00094251| 2.68| 4.53|\n\nAs observed in our experiments (Table -1), the attention values do not strictly follow an exponential pattern of the form $3^i$. While our construction uses attention weights proportional to $3^i$ for illustrative purposes, we would like to emphasize that the proof of our representation result relies only on a much weaker condition. Specifically, for a token at position $n$, it suffices that the attention weight at position $n - k$, denoted $att_{n, n-k}$, is greater than both:\n\n1. The sum of the attention weights from positions $n - 1$ to $n - k + 1$, and  \n2. All other attention weights.\n\nTo elaborate further, in the first layer of attention, we obtain: $v_{n} = \\sum_{i=1}^{\\min(n, k)}  att_{n, n-i} e^x_{n-i}$ where $att_{n, n-i}$ is the attention weight at position $n - i$, when the token at position $n$ is in consideration, $k$ refers to the order of the markov chain, and $e^x_{n-i}$ is the value vector (one hot vector of the symbol associated with $x_{n-i}$ by construction) for token $x_{n-i}$. The MLP is then used to compute $u_n = \\sum_{i=0}^{\\min(n, k-1)} att_{n, n-i} e^x_{n-i}$ by isolating $e^x_{n-k}$ from $v_n$, subtracting it, and subsequently adding $e^x_{n}$. This is achieved via a weighted linear combination using the MLP and skip connections. Specifically, $e^x_{n}$ is introduced through skip connections, while the MLP plays a central role in isolating and removing $e^x_{n-k}$ from $v_n$ with an appropriate weight, which helps in obtaining $u_n$. Therefore, the key condition required for this construction to function, as also noted in the supplementary material (lines 555–560), is: $att_{n, n-k} > \\sum_{i=1}^{k-1} att_{n, n-i} $ which guarantees that $e^{x}_{n-k}$ can be effectively isolated. Our experiments (Table -1) confirm that this condition holds consistently across different settings. We will update the paper to include these additional experimental results, as well as an explicit statement of the weaker, but sufficient, theoretical condition.\n\n**Q3. In Appendix E (fig 5/9/13), why not compute a pseudo-attention map for the first layer?**\n\nThe pseudo-attention map represents the idealized attention pattern required to approximate a conditional $k$-gram estimator. In other words, it highlights the tokens that need to be attended to in order to estimate the transition probabilities of a $k$-th order Markov chain, effectively functioning as a $k$-th order induction head. Please refer to Figure 1b (page 3) of the main paper for an illustration. While it is possible to compute the pseudo-attention map for the first layer as well, we note that the role of the first-layer head is primarily to copy the previous $k$ tokens (i.e., from positions $n-1$ to $n-k$) for the $n$-th token. For this reason, we believe that displaying the pseudo-attention map for the first layer may not offer significant additional insight. However, a potentially useful comparison would be to show the attention map of our theoretical construction in the first layer alongside the learned attention map. We agree that this could help clarify how closely the learned mechanism approximates the idealized behavior. We will include this plot in the final version of the paper.\n\n**Q4. In Fig 2 Layer 2 (likewise figs 6-15) -- is it correct ...  a more interesting plot?**\n\n It is indeed correct that the pseudo attention is ill-defined for the first k tokens, and accordingly, these tokens must be excluded when comparing it with the learned attention. Furthermore, as the sequence length increases, the estimation becomes more accurate. We also greatly appreciate the thoughtful suggestion to consider sequences of length $k + N$ and to present the attention matrix for a window of length $N$ starting from the $(k + 1)$-th token. We will incorporate this into revised version of the paper.\n\n**Q5. Have you considered the implications for inference ... against LLMs for more circumstantial evidence on the representations within LLMs.**\n\nThis is indeed an interesting direction, and while our primary focus was on deriving the tightest characterization of the relationship between transformer depth and the Markov order required to solve the Markov in-context learning task, we agree that exploring quantized inference could offer valuable insights. For additional context, our representational results adopt the formulation from [1], which shows that $\\Omega(\\log T + k)$ bits per parameter are sufficient to achieve an additive approximation error of $O(1/T)$, where $T$ represents the sequence length and $k$ is the Markov order. That said, we did conduct a preliminary analysis of weight and activation quantization on a trained model using a second-order Markov chain with sequence length 32. We computed the Frobenius norm between the pseudo attention map (optimal attention map) and the attention map of the second layer of the quantized model across 10 random sequences, and report the mean and variance below:\n\n| Quantization (bits) | Frobenius Norm (Mean ± Std) |\n|-|-|\n| 2                   | 4.6108 ± 0.7190             |\n| 4                   | 4.5842 ± 0.9826             |\n| 8                   | 3.8615 ± 0.4590             |\n| 32 (No quantization)| 3.7700 ± 0.4725 |\n\nThese results suggest that 8-bit quantization yields performance that is quite close to full-precision (32-bit), at least in terms of preserving the attention structure as captured by the Frobenius norm, which is indeed an interesting implication and also seems to match the theoretical bounds. Due to the updated NeurIPS policy, we are unable to share the attention maps directly. However, we will include these quantization results as well as the attention maps and their implications in the revised version of the paper.\n\n---\n\n**References**\n\n[1] Rajaraman, Nived, et al. \"Transformers on markov data: Constant depth suffices.\" NeurIPS 37"}, {"review_id": "sbx0YRF25h", "reviewer": "Reviewer_hmwe", "summary": "This paper aims to prove that a transformer with two layers within which contains a single head self-attention is sufficient to realise the k-th order in-context learning for any k. This is so far the tightest (2 layers and 1 head in each layer) theoretical solution against previous works where either more layers or more heads are needed.", "questions": "As suggested above, can you provide a more friendly description on how the MLP weights are constructed?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n- This paper provides a rigorous proof to the requirements of induction heads in the transformer architecture.\n- This paper highlights the importance of the MLP module in a transformer layer. With a proper construction of MLP weights, the tightest results are obtained. Whereas previous works which requires more layers just ignored MLPs. This is also an important evidence suggesting that the MLPs in the pratical LLMs can play roles in structural abilities other than merely Key-Value stores.\n- Empirical results in a simplified 1st-order markov analysis have accorded with the theorical prediction by a large margin.\n\nWeaknesses\n\n- Despite the simplification of the experiment as stated by the authors, the analysis still relies on synthesized data training other than induction heads in real world problems. The gaps are underexplored.\n- In order to convey a concise and clear thought, this paper sacrifices much details. While most of them are acceptable, how the 1st layer MLP constructs $u_n$ is not easy to understand. Since this part is the key difference of this paper's constructional proof, it may worth more spaces (spared from the warm-up section) to describe (or even with more figures).", "was_revised": true, "final_justification": "I’ve read the authors’ rebuttal and rebuttals to other reviews. The concerns have been addressed. Through this is a theoretical work and its setting is significantly simpler to real-world LLM corpus, it reveals the functions of both attentions and FFNs are interconnected on a common transformer structure. Therefore I would consider maintaining the score", "rebuttal": "We sincerely thank the reviewer for their thoughtful comments, valuable suggestions, and recognition of our work. Below, we provide detailed responses to each of the individual comments and questions.\n\n---\n\n## Comments on weaknesses\n\n**W2. In order to convey a concise and clear thought, this paper sacrifices much details. While most of them are acceptable, how the 1st layer MLP constructs $u_n$ is not easy to understand. Since this part is the key difference of this paper's constructional proof, it may worth more spaces (spared from the warm-up section) to describe (or even with more figures).**\n\nWe respectfully apologize for not elucidating the result more clearly. Below, we try to provide a more accessible description of how the MLP functions in our construction. We will also make sure to elaborate on this in the revised version of the paper to improve the readability of this section.\n\n---\n\n## Questions\n\n**Q1. As suggested above, can you provide a more friendly description on how the MLP weights are constructed?**\n\nWe clarify below how the MLP is used in our construction. We first provide a brief sketch of what the function of the MLP is and then dive into the details.\n\n**Brief Sketch:** In the first layer of attention, we obtain: $v_{n} = \\sum_{i=1}^{\\min(n, k)}  att_{n, n-i} e^x_{n-i}$ where $att_{n, n-i}$ is the attention weight at position $n - i$, when the token at position $n$ is in consideration, $k$ refers to the order of the markov chain, and $e^x_{n-i}$ is the value vector (one hot vector of the symbol associated with $x_{n-i}$ by construction) for token $x_{n-i}$. The MLP is then used to compute $u_n = \\sum_{i=0}^{\\min(n, k-1)} att_{n, n-i} e^x_{n-i}$ by isolating $e^x_{n-k}$ from $v_n$, subtracting it, and subsequently adding $e^x_{n}$. This is achieved via a weighted linear combination using the MLP and skip connections. Specifically, $e^x_{n}$ is introduced through skip connections, while the MLP plays a central role in isolating and removing $e^x_{n-k}$ with an appropriate weight.\n\n\n**More Details:** The first layer of the MLP is designed to extract the identity of the symbol that appeared $k$ steps before the current position $n$, i.e., to identify $x_{n-k}$. Recall that the attention mechanism in the previous stage computes: $v_n = \\sum_{i=1}^{\\min(n, k)}att_{n, n-i} e^x_{n-i}$ where each $att_{n, n-i} = \\frac{3^i}{C_{attn}}$ (where $C_{attn}$ is the normalization constant of attention i.e. $C_{attn}=\\sum_{i=1}^k 3^i$), and $e^x_{n-i}$ is the value vector for token $x_{n-i}$, represented as a one-hot vector by construction. Because the attention weights grow exponentially with $i$, the contribution from $x_{n-k}$ is strictly larger than the total contribution from all more recent symbols. That is: \n\n$att_{n, n-k} = \\frac{3^k}{C_{attn}} > \\sum_{i=1}^{k-1} att_{n, n-i} = \\sum_{i=1}^{k-1} \\frac{3^i}{C_{\\texttt{attn}}} = \\frac{3^{k} - 3}{2C_{\\texttt{attn}}}$ ... (Inequality - 1)\n\nThis inequality ensures that the coordinate in $v_n$ corresponding to $x_{n-k}$ is strictly larger than any combination of more recent contributions. \n\n*Layer-1:* To isolate $x_{n-k}$, the first MLP layer contains the one hot encodings corresponding to the vocabulary and hence one can take a dot product with this to $v_n$. After taking the dot product, the coefficients associated with each symbol can be distinctly identified. Since the vocabulary is embedded in the first layer, these coefficients naturally reside in orthogonal subspaces. The bias of the first layer is designed in a specific way to be able to isolate the symbol $x_{n-k}$ and this is done  by using the result shown in Inequality-1). Hence, the bias is a constant vector where each dimension has the same constant: $-\\frac{3^{k} - 3}{2C_{attn}}$.  After adding this bias, the components corresponding to $x_{n-1}, x_{n-2}, \\ldots$ become negative while the component for $x_{n-k}$ remains positive. Applying ReLU zeroes out the negative parts, leaving only the positive coordinate, corresponding to $x_{n-k}$, and after passing it through the norm layer, we obtain a one hot vector that that indicates the identity of $x_{n-k}$ (which is basically $e^x_{n-k}$ as its a one hot vector). Note that the first layer of the MLP is constructed in such a way that $e^x_{n-k}$ and $v_n$ are present in orthogonal subspaces. \n\n*Layer-2:* Since, the first layer extracts the one hot vector that indicates the identity of $x_{n-k}$, the second layer is then used to obtain $u_n$. The weights of the second layer are then designed to not only multiply  $e^x_{n-k}$, $v_n$ as well as $e^x_{n}$ (which is obtained via skip connections) with appropriate weights, but also to extract and aggregate them within the same subspace to obtain $u_n$. This is done by, designing the second layer wieghts of the MLP, that multiply $v_n$ with $3$ and $e^x_{n-k}$ with $-3^k/C_{attn}$ and finally $e^x_{n}$ with $1/C_{attn}$. Hence, once the weighted aggregation is computed, one can observe that \n\n$u_n = \\frac{1}{C_{attn}} e^x_{n}  + 3 v_n -\\frac{3^k}{C_{attn}} e^x_{n-k}$\n\nThus, the first  two layers of the MLP layer performs the task of obtaining $u_n$. We will revise the manuscript to include a more accessible explanation in the main text."}, {"review_id": "Wsa6JtG55L", "reviewer": "Reviewer_GjjU", "summary": "The paper shows that a two-layer transformer with a single attention head per layer can represent the conditional k-gram model for arbitrary Markov order k (Thm 4.1) by using the universal approximation of MLP. It further provides \n- (i) a depth-width trade-off construction with two heads in the first layer (Thm 3.2) \n-  (ii) a gradient-descent analysis showing that a simplified two-layer network learns first-order induction heads (Thm 5.1). \n\nEmpirical evidence is limited to visualizing attention heatmaps on synthetic data (Fig. 2).\n\n\nWhile the authors successfully construct a representation of the  $k$-gram model using a two-layer Transformer with a single attention head per layer, the construction heavily relies on the universal approximation capability of the MLP, which likely introduces a substantial number of parameters. This reliance renders the comparison with alternative architectures—such as deeper models or those with more attention heads but less dependence on the MLP—potentially unfair. Addressing these empirical gaps, along with providing a fair and comprehensive comparison of the total parameter counts across models, would substantially enhance the contribution.", "questions": "- Could the authors provide a simplified explanation of the key technical challenges that arise when extending the optimization analysis to the case $k>1$. \n- The reviewer is curious whether the hand-crafted attention pattern proposed in the construction actually emerges in practice. Could the authors present attention scores from real-world tasks to examine whether such patterns are observed empirically?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n-  Prior work has demonstrated that representing the conditional $k$-gram model typically requires either a 3-layer Transformer with a single attention head per layer or a 2-layer Transformer with two heads per layer. The proposed new construction only use a 2-layer Transformer with a single attention head per layer. Therefore, they answer the question \"can a 2-layer, 1-head transformer represent the conditional k-gram\" successfully, bridging the gap between $1$-layer and $3$-layer.\n- The paper constructs the representation using the non-non-linearities of MLP. The proof sketch highlights how ReLU + LayerNorm can simulate an additional “virtual head”, a perspective likely useful beyond this setting. \n\n\n**Weakness**\n\n-  The proposed construction necessitates that the MLP in the first layer effectively encodes the entire one-hot vector within its weight matrix, which likely entails a substantial number of parameters. However, the authors appear to account only for the parameter counts in the embedding layer and the query, key, and value matrices, while omitting those in the MLP. As a result, the parameter comparison with alternative architectures may be misleading. For a fair and comprehensive evaluation, the total parameter count should include all components—namely, the embedding layer, attention mechanism, LayerNorm, and MLP—for each architecture under consideration.\n-  The assumptions underlying the gradient descent analysis appear highly artificial—for instance, Assumption (iv) involves a carefully designed initialization coupled with frozen, \"optimal\" MLP blocks. Consequently, it remains unclear whether standard stochastic gradient descent applied to a genuine two-layer, single-head transformer would converge to the proposed representation of the $k$-gram model. To substantiate the theoretical claims, it is recommended to conduct empirical experiments using the actual two-layer single-head architecture and to examine whether the trained MLP replicates the functional behavior of encoding the entire one-hot vector, as prescribed in the hand-crafted construction.\n- The paper lacks empirical comparisons between the two-layer single-head architecture and the three-layer single-head counterpart on realistic datasets such as natural language, music, or other real-world sequences that can be effectively modeled by the $k$-gram framework.", "was_revised": true, "final_justification": "The authors have addressed some of my concerns. As a result, I am raising my score to 4.", "rebuttal": "We thank the reviewer for the helpful feedback and insightful comments. We address the individual comments and questions below. We respectfully believe that there maybe some misunderstandings about the contributions of our work and also try to clarify them.\n\n---\n## Clarifications on potential misunderstandings\n\n**C1. The paper shows ... universal approximation of MLP. (Universal approximation of MLP)**\n\nWe would like to clarify that **our representation result does not rely on the universal approximation theorem of an MLP**. Instead, we provide an **explicit and constructive design** that uses an MLP of depth 3, where each linear layer is a finite-dimensional matrix of size $(6S + 3) \\times (6S + 3)$, with $S$ denoting the vocabulary size. We include a high-level sketch of the MLP's function here (please see Section 4, lines 210–222 in the paper, and the appendix for full details). The first attention layer computes a geometrically weighted average over the previous $k$ tokens, yielding a vector $v_n$. The MLP then reconstructs a second vector $u_n$, resembling a hard attention over the current and preceding $k-1$ tokens, without using an additional attention head. This is achieved via careful weight encoding, bias selection, ReLU activations, and normalization to algebraically extract and combine embeddings (notably $x_n$, $x_{n-k}$, and $v_n$) using orthogonal subspaces and skip connections. \n\n**C2. While the authors ... would substantially enhance the contribution. (Parameter Count)**\n\nWe would like to reiterate that our representation results **do not rely on the universal approximation theorem**. Instead, we provide an explicit construction that uses a 3-layer MLP, where each linear layer is a finite-dimensional matrix of size $(6S + 3) \\times (6S + 3)$, with $S$ denoting the vocabulary size. Importantly, the **parameter count reported in Table 1 in the paper includes the full MLP used in our construction**. This ensures a fair and accurate comparison with prior works. We would also like to point out that Rajaraman et al. [1] use three separate 2-layer MLPs, each with linear layers of the same dimensions ($(6S + 3) \\times (6S + 3)$). This design choice is similarly reflected in the parameter count for their model as shown in Table-1 in the paper. We would like to note that the parameter counts in the paper (Table-1) account only for the most important components in both our construction and that of Rajaraman et al. [1], namely the query, key, and value layers across different attention layers, the weights of the MLPs, and the relevant positional encodings. Below, we provide an explicit comparison that takes all components into account. We show that our model is not only **still more parameter-efficient than Rajaraman et al. [1], but the margin is even larger than initially reported**. Note that even in this setting, our construction is still more efficient when compared to Nichani et al. [2] as well. To clearly illustrate this, we now provide the full parameter breakdown for both our method and Rajaraman et al. [1], along with a comparative table summarizing prior constructions for a $k$-th order Markov chain. Here, $S$ denotes the vocabulary size and $T$ the sequence length.\n\n**Our Construction (Total: $9(6S + 3)^2 + (6S + 3)(2T + 2S + 9)$)**\n\n1. Token Embedding: $(6S + 3) \\times S $\n2. Attention Layer 1:  \n   a. Positional Encoding: $(6S + 3) \\times T$  \n   b. Q, K, V: $3(6S + 3)^2$\n3. MLP (after Attn 1):  \n   a. Weights: $3(6S + 3)^2$  \n   b. Biases: $3(6S + 3)$  \n   c. Layer Norms: $6(6S + 3)$\n4. Attention Layer 2:  \n   a. Positional Encoding: $(6S + 3) \\times T$  \n   b. Q, K, V: $3(6S + 3)^2$\n5. Output Projection: $(6S + 3) \\times S$\n\n**Rajaraman et al. [1] Construction (Total: $15(6S + 3)^2 + (6S + 3)(3T + 2S + 10)$**\n\n1. Token Embedding: $(6S + 3) \\times S$\n2. Transformer Layer 1:  \n   a. Positional Encoding: $(6S + 3) \\times T$  \n   b. Q, K, V: $3(6S + 3)^2$  \n   c. MLP: $2(6S + 3)^2 + 2(6S + 3)$  \n   d. Layer Norm: $2(6S + 3)$\n3. Transformer Layer 2: (same as Layer 1)\n4. Transformer Layer 3:  \n   a. Positional Encoding: $(6S + 3) \\times T$  \n   b. Q, K, V: $3(6S + 3)^2$  \n   c. MLP: $2(6S + 3)^2 + 2(6S + 3)$\n5. Output Projection: $(6S + 3) \\times S$\n\nHence finally the table for all the kth order construction is:\n\n\n| **Previous works**| **Markov order** | **Architecture**| **# Parameters**|\n|-|-|-|-\n| Nichani et al. (2024) [2] | $k$| $2$-layer, $k$-head| $(k+(k+1)^2)(S+T)^2 + T(S+T)$|\n| Rajaraman et al. (2024) [1] | $k$| $3$-layer, $1$-head| $15(6S + 3)^2 +  (6S + 3)(3T + 2S +10)$|\n| **Ours**|$k$| $2$-layer, $1$-head| $9(6S + 3)^2 + (6S + 3)(2T + 2S + 9)$|\n\nWe will include these results in the revised version of the paper.\n\n---\n## Comments about weaknesses\n\n**W1. The proposed construction...for each architecture under consideration.**\n\nPlease refer to the second clarification.\n\n**W2. The assumptions underlying  ... in the hand-crafted construction..**\n\nWe believe there may have been a misunderstanding. **All experimental results in the paper (Fig. 2 and Figs. 5–16) were obtained using standard initialization schemes, without assuming any component to be optimal**. These results demonstrate that, under standard initialization, a two-layer, single-head transformer converges to a $k$-gram estimator. Our theoretical construction is consistent with the empirical results, as described in lines 205 to 209 of the paper. Given the highly non-convex nature of the problem, there may be multiple valid constructions showing how two-layer, single-head transformers can represent induction heads. Our approach is one such construction. To further support this, we trained two models: (1) a standard two-layer, single-head transformer, and (2) a modified version where the vocabulary is embedded into the MLP, in line with our theoretical construction, and data generated from second order markov chains. We ran each model three times and computed the difference of their cross-entropy loss to the cross-entropy obtained using the true transition probabilities over a large data batch. The results are closely aligned. Additionally, the attention maps for both models  of the second layer closely match the ideal attention patterns of a $k$-gram estimator. We will include these results in the final version of the paper. We would also like to emphasize that while our construction demonstrates a specific way in which induction heads can be represented, it is not obvious that gradient descent will necessarily find this structure. Given the problem's non-convexity, multiple solutions are possible. Our main contribution is a representation result, and the fact that gradient descent empirically converges to a solution that closely matches our construction suggests that the relevant structure can emerge naturally during training.\n\n| Embedded Vocabulary     | Full Model (Standard Initialization) | Difference in Mean |\n|-|-|-|\n| 0.1004 ± 0.00128| 0.100223 ± 0.00110| 0.00017|\n\n**W3. The paper lacks empirical comparisons ..  real-world sequences that can be effectively modeled by the k-gram framework.**\n\nWe would like to note that our work is primarily theoretical in nature, with the goal of obtaining the tightest possible characterization of the relationship between transformer depth and Markov order in the context of in-context learning for Markov processes. We acknowledge that the low-depth settings we study may not directly reflect practical applications, as also suggested in prior work ([3]), where larger depths are required to solve realistic tasks. Nonetheless, we believe that these simplified settings provide valuable insights into the fundamental mechanisms by which transformers can implement learning algorithms such as $k$-gram estimation. Extending these results to more realistic settings is an important direction for future work.\n\n--- \n\n## Questions\n\n**Q1. Could the authors provide ... optimization analysis to the case $k>1$?**\n\nAs noted in lines 39–47 in the paper, prior works have also found optimization analysis for $k > 1$ to be challenging. [2] provides such analysis only for first-order Markov chains and on a specialized disentangled transformer. [1] gives a representation result for realistic architectures (with MLPs and LayerNorms), but no optimization analysis. Our work provides the tightest known construction relating transformer depth to Markov order under realistic assumptions and also initiates an optimization analysis in this setting. As Reviewer VB7N also observed, the problem is inherently difficult and progress requires reasonable assumptions. The main obstacles in extending to $k > 1$ is that the MLP contains layer norms and also that MLP's optimality depends on $v_n$ (the output from the first layer) . Due to this, the gradient descent now involves two things, the first being the coupled relationship of the first layer and the MLP and the second being performing optimization analysis with layer norms (which is shown to be difficult to do in [4]). Together, due to the coupled behavior as well as dependency on layer norms, makes the analysis challenging. While a full analysis for $k > 1$ may be possible, we leave this important direction to future work.\n\n**Q2. The reviewer ...  empirically?**\n\nDue to space constraints, we respectfully refer you to our responses to Questions 1 and 2 from Reviewer 2Q6Z, where we have tried to address a closely related question in detail and we hope our prior responses may also be helpful in this context. However, we would be happy to elaborate further if additional clarification is needed.\n\n---\n**References**:\n\n[1] Rajaraman, Nived, et al. \"Transformers on markov data: Constant depth suffices.\" NeurIPS 37\n\n[2] Nichani, Eshaan et. al. \"How transformers learn causal structure with gradient descent.\"ICML 2024.\n\n[3] Zhao, Wayne Xin, et al. \"A survey of large language models.\" (2023).\n\n[4] Xu, Jingjing, et al. \"Understanding and improving layer normalization.\" NeurIPS 32."}, {"review_id": "xEqpEzs6hR", "reviewer": "Reviewer_VB7N", "summary": "This is a theory paper on multi-layer Transformer models. The paper studies inductions heads for the conditional k-grams problem: given a sequence of k tokens, can the model predict the next token correctly?  This problem is modeled as a Markov predictions in a k-gram and has connections with k-th order Markov chains and In Context Learning. \nAuthors establish that a 2-layer Transformer model is sufficient for building an induction head for this problem, regardless of the value of k. This is an improvement on previously known result which had established that 3 layers are sufficient. Authors also share a result on the dynamics of gradient descent for converging to the solution.", "questions": "1. optimization: strong assumptions are made. What happens if this is not the case that all layers have their optimal values? how far are we from this setup when data is noisy?\n\n2. are 2 layers the fewest needed layers? how hard is it to establish a lower-bound?", "limitations": "more numerical experiments on noisy / real data could help bridge the gap between theory results and practical architecture selection", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. the paper is well written and easy to follow\n2. the result is well exposed: reduction of 3 to 2 layers \n3. literature review is complete and nicely exposed\n\nWeaknesses\n\n1. optimization result makes very strong assumptions. I understand that the problem is really hard and unless we make strong assumptions it is not possible to prove anything useful. However we also need to understand the behavior outside the nice regime proposed here. Authors are assuming all layers but one already have the optimal value. \n\n2. lack of lower bound for number of layers\n\n3. optimization is disconnected from realistic situations because very strong assumptions are needed.\n\n4. limited numerical experiments: gap between theory and practice", "was_revised": false, "final_justification": "", "rebuttal": "We would like to sincerely thank the reviewer for their thoughtful feedback and insightful comments. We address the individual comments and questions below.\n\n---\n\n## Comments about weaknesses\n\n**W1.  optimization result makes very strong assumptions. I understand that the problem is really hard and unless we make strong assumptions it is not possible to prove anything useful. However we also need to understand the behavior outside the nice regime proposed here. Authors are assuming all layers but one already have the optimal value.**\n\nWe believe this question can be interpreted in two ways: one theoretical and the other practical. We address both perspectives separately.\n\n*Theoretical assumptions*: Prior work has shown that full optimization analysis in this setting remains a challenging problem. Nichani et. al. [2], for example, provides such analysis only for first-order Markov chains and a specialized disentangled transformer. Rajaraman et. al. [1] offers a representation result for more realistic architectures (including MLPs and LayerNorms), but does not address optimization. In contrast, our work presents the tightest known construction relating transformer depth to Markov order under realistic architectural assumptions and initiates an optimization analysis in this setting. It is also worth noting that Nichani et. al. [2] makes comparable assumptions by analyzing optimization on a reduced model of a disentangled transformer, where the second attention layer relies on a single scalar, an assumption similar to ours. We agree with the reviewer that a full optimization analysis remains an important open question. Our primary focus was on the representational results, while also making partial progress on learning dynamics. A complete characterization is left for future work.\n\n*Practical Experiments:* We would like to emphasize that all experiments in our paper were conducted using a two-layer, single-head transformer model with standard initialization schemes. No component was assumed to be optimal. Further details are provided in our responses to Weaknesses 3 and 4.\n\n**W2. lack of lower bound for number of layers**\n\nWe would like to clarify that, as discussed in lines 39–47 of the paper, Sanford et al. [3] have provided a formal lower bound in their work titled **\"One-layer transformers fail to solve the induction heads task.\"** Using a communication complexity argument, they demonstrate that no one-layer transformer can solve the induction heads task unless its size is exponentially larger than that of a two-layer transformer capable of solving the task. This result establishes a meaningful lower bound on the number of layers required.\n\n**W3. optimization is disconnected from realistic situations because very strong assumptions are needed.**\n\nWe kindly refer the reviewer to our response to the first weakness for additional context. Importantly, all experimental results presented in the paper (Fig. 2 and Figs. 5–16) were obtained using standard initialization schemes, without assuming any component to be optimal. These experiments show that, under standard initialization, a two-layer, single-head transformer reliably converges to a $k$-gram estimator. Moreover, our theoretical construction aligns closely with these empirical findings, as detailed in lines 205–209 of the paper.\n\n**W4. limited numerical experiments: gap between theory and practice**\n\nWe would like to emphasize that all experimental results in the paper (Fig. 2 and Figs. 5–16) were obtained using standard initialization schemes, without assuming any component to be optimal. These findings demonstrate that, under such initialization, a two-layer, single-head transformer converges to a $k$-gram estimator. Additionally, our theoretical construction closely aligns with the empirical observations. We also perform experiments on noisy data and we present our results below (please see the reply to question -1).\n\n---\n\n## Questions\n\n**Q1. optimization: strong assumptions are made. What happens if this is not the case that all layers have their optimal values? how far are we from this setup when data is noisy?**\n\nWe believe this question can be interpreted in two ways: one theoretical and the other practical. We believe that we have attempted to address both aspects in our responses to Weaknesses 1, 3, and 4. Additionally, below we provide experimental results for our model trained on noisy data.\n\n**Noisy Data:** To evaluate robustness to noise, we sampled sequences from second-order Markov chains and introduced noise with probability $p_n$: for each token, with probability $p_n$, we replaced it with a random token; otherwise, it was left unchanged. Each model was run three times, and we compared their cross-entropy loss to the cross-entropy computed using the true transition probabilities over a large batch. We observe that when $p_n = 0$ (no noise), the ideal attention map closely aligns with the learned attention map in the second layer. However, as $p_n$ increases, this alignment degrades. Notably, even at $p_n = 0.1$, the pseudo attention map already starts to differ significantly. Additionally, the gap between the model's cross-entropy loss and the cross entropy loss computed using the true transition probabilities steadily increases with higher noise levels:\n\n| $p_n$ | Difference in Cross-Entropy Loss|\n|-|-|\n| 0| $0.10 \\pm 0.003$|\n| 0.1| $0.15 \\pm 0.010$ |\n| 0.25| $0.18 \\pm 0.002$|\n| 0.5|$0.21 \\pm 0.005$|\n\nWe emphasize that our theoretical construction is designed for estimating $k$-th order Markov chains in the absence of noise. We agree with the reviewer that understanding the impact of noise on such constructions is an interesting direction, and we leave a thorough theoretical treatment under noisy conditions to future work. We will add these results to the revised version of the paper. \n\n\n**Q2. are 2 layers the fewest needed layers? how hard is it to establish a lower-bound?**\n\nAs described in the response to Weakness-2, Sanford et. al. [3] already establishes a lower bound. Therefore, we believe that two layers are the minimal number required.\n\n---\n\n**References:**\n\n[1] Rajaraman, Nived, et al. \"Transformers on markov data: Constant depth suffices.\" NeurIPS 37\n\n[2] Nichani, Eshaan et. al. \"How transformers learn causal structure with gradient descent.\"ICML 2024.\n\n[3] Sanford, Clayton, et. al. \"One-layer transformers fail to solve the induction heads task.\" arXiv preprint arXiv:2408.14332 (2024)."}]

Author rebuttal:

Meta-review: This work showed that a two-layer transformer with a single head can realise any conditional k-gram for in-context learning for any k, which is so far the tightest theoretical construction. It also provided experimental results that support the theoretical analysis. It also explores the feasibility of discovery by training dynamics. 

The discussions have helped clarify/address a few concerns from the reviewers, e.g., clarifying the experimental setup, adding experiments. While the setting is simpler than real-world scenarios, the novel and tight theoretical construction and the empirical support are solid contributions to the community. Please incorporate the clarifications/revisions/added experiments into the final version if accepted.
Decision: Accept (spotlight)