Selective Induction Heads: How Transformers Select Causal Structures in Context
Authors: Francesco D'Angelo
Francesco Croce
Nicolas Flammarion
Venue: ICLR 2025 Poster (2025)
Keywords: Transformers
Markov chain
interpretability
attention
in-context learning
Primary area: foundation or frontier models, including LLMs

Abstract: Transformers have exhibited exceptional capabilities in sequence modelling tasks, leveraging self-attention and in-context learning. Critical to this success are induction heads, attention circuits that enable copying tokens based on their previous occurrences. In this work, we introduce a novel synthetic framework designed to enable the theoretical analysis of transformers’ ability to dynamically handle causal structures. Existing works rely on Markov Chains to study the formation of induction heads, revealing how transformers capture causal dependencies and learn transition probabilities in-context. However, they rely on a fixed causal structure that fails to capture the complexity of natural languages, where the relationship between tokens dynamically changes with context.  To this end, our framework varies the causal structure through interleaved Markov chains with different lags while keeping the transition probabilities fixed. This setting unveils the formation of Selective Induction Heads, a new circuit that endows transformers with the ability to select the correct causal structure in-context. We empirically demonstrate that attention-only transformers learn this mechanism to predict the next token by identifying the correct lag and copying the corresponding token from the past. We provide a detailed construction of a 3-layer transformer to implement the selective induction head, and a theoretical analysis proving that this mechanism asymptotically converges to the maximum likelihood solution. Our findings advance the theoretical understanding of how transformers select causal structures, providing new insights into their functioning and interpretability.

Published as a conference paper at ICLR 2025
SELECTIVE INDUCTION HEADS: How TRANSFORM- ERS SELECT CAUSAL STRUCTURES IN CONTEXT
Francesco D' D'Angelo, Francesco Croce & Nicolas Flammarion Theory of Machine Learning Lab, EPFL., Lausanne, Switzerland
ABSTRACT
Transformers have exhibited exceptional capabilities in sequence modelling tasks, leveraging self-attention and in-context learning. Critical to this success are induc- tion heads, attention circuits that enable copying tokens based their previous currences. this work, we introduce a novel synthetic framework designed to en- able the theoretical analysis of transformers' ability dynamically handle causal structures. Existing works rely on Markov Chains to study the formation of in- duction heads, revealing how transformers capture causal dependencies and learn transition probabilities in-context. However, they rely on fixed causal structure that fails to capture the complexity of natural languages, where the relationship between tokens dynamically changes with context. To this end, our framework varies the causal structure through interleaved Markov chains with different lags while keeping the transition probabilities fixed. This setting unveils the formation Selective Induction Heads, a new circuit that endows transformers with the abil- ity to select the correct causal structure in-context. We empirically demonstrate that attention-only transformers learn this mechanism to predict the next token by identifying the correct lag and copying the corresponding token from the past. We provide a detailed construction of a 3-layer transformer to implement the selective induction head, and a theoretical analysis proving that this mechanism asymptot- ically converges to the maximum likelihood solution. Our findings advance the theoretical understanding of how transformers select causal structures, providing new insights into their functioning and interpretability.
INTRODUCTION
As autoregressive generative models continue to scale and are increasingly deployed in real-world applications, the question of how Transformer models ( (Vaswani et al., 2017) function internally becomes pressing. Yet the inherent complexity of natural language hinders the ability to fully com- prehend how these models make decisions and work internally. To address this challenge, many recent works have attempted to formulate synthetic frameworks that simplify the problem and en- able theoretical analysis while still capturing the remarkable properties and phenomena observed in large language models, such as in-context learning (Brown, 2020; Garg et al., 2022: Bai et al., 2024; Von Oswald et al., 2023a; Sander al., 2024). Mechanistic interpretability (Olsson et al., 2022) emerges as a line of research focused on reverse-engineeringg the complex computations per- formed inside transformer in order understand how a certain output produced for given input. This research has uncovered the formation of induction heads (Olsson et al., 2022) i.c., in- terpretable circuits embedded within the transformer's weights, capable of simple operations such as copying tokens. By examining such circuits and their combinations, one can understand the al- gorithms that transformers implement to solve a given task. For instance, (Nichani al., 2024; Bietti et al., 2024; Edelman et al., 2024) demonstrated that induction heads enable transformers to implement in-context bigrams for next-token predictions in Markov Chains. Such mechanisms are not limited to simplified models: Nguyen (2024) showed that transformers may rely N- N-gram rules even in natural language processing. Yet the process by which transformers select between such learned rules remains poorly understood. While in-context learning studies how transformers solve tasks from demonstrations in the prompt, in-context selection focuses on how transformers
"Corresponding authors: francesco.dangelo@epp.coo

---

Published as a conference paper at ICLR 2025
interleaved lag
Task: prodict the
the correct
structime
with
How do 3-layer attention-only transformers solve this task?
Laper
Attatie
Stterbor Laye bog
LI: extrazes single transition probahilities each lig
122: ap aggregates tmassition probubilities from the posi
L3: selective Induction head select most likely
Figure 1: Summary of the framework. Top: we define a new task based on Interlcaved Markov Chains of different lags (k = and & = 2 in the example). Middle: given a sequence generated from chain of unknown lag, the model has identify the true lag, and use it to predict the distribution of the next token. Bottom: attention-only transformers can solve this task with layers. The first computes the transition probabilities for each lag seen during training, the second aggregates these probabilities over the entire past, and finally third layer implements the selective induction head, which selects the correct lag.
select the most suitable approach to solve a given task from those encountered during training, using instances present in the context. For example, Bai et al. (2024) examines how transformers per- form in-context algorithm selection when pre-trained on mixtures linear and logistic regression with different noise. Similarly, Yadlowsky al. (2023) studied the ability of transformers to per- form model selection between different function class families. Our work takes another step toward understanding this selection process, leveraging Markov chains with different causal structures.
In-context causal structure selection. Recently, Markov chains have been employed to formu- late interesting to-rencee---eeueeee tasks that can be solved by transformers with interpretable solutions (lidiz et al., 2024; Nichani et al., 2024; Makkuva et al., 2024; Edelman et al., 2024). In particular, Nichani et al. (2024) show that transformers trained on Markov Chain sequences learn circuits that capture the causal structure, i.e., the set of parent tokens for each token in the sequence and estimate transition probabilities in-context. The existing works relying on Markov chains, fail to model the nuanced relationships typical of natural language. The same word pair can have dif- ferent causal relationships depending on the surrounding context. While an effective model should recognize these contextual dependencies, previous research has overlooked this consideration by adopting fixed causal structures. To address this limitation, we propose a new synthetic lask de- signed to mimic different causal dependencies (Sec. 3). We consider Interleaved Markov Chains, with fixed transition probabilities between states but different underlying causal structures (Fig, 1), and theoretically study how 3-layers attention-only transformers learn correctly predict the next token in sequence.
Selective induction heads. To solve the task hand-correctly predicting the next token in-context in sequence generated within this setup (middle of Fig. 3)--rrnnforrmrss need to leam circuit that adapts to the given context select the correct causal structure among those seen during training. We call this circuit selective induction head, as it differs fundamentally from the induction heads introduced so far the literature, where the circuit learns either to copy token from a certain position fixed by the unique structure of the data by comparing its semantics. In our task, the transformer (with attention maps depicted Fig. 1) needs to learn to aggregate all past information to determine from which past position the corresponding token should be copied in order predict the next token.
A transformer construction for in-context selection. To understand and formalize the selective heads, we provide an interpretable construction of the self-attention layer weights in a 3-layer attention-only disentangled transformer (Friedman al., 2023) that implements this mechanism
2

---

Published as a conference paper at ICLR 2025
(Sec. 4). We empirically demonstrate that the constructed transformer matches the performance of both disentangled and standard transformed trained from scratch (Sec. 5) and that 2-layer attention- only transformers cannot solve the task. Moreover, observe that the attention maps of the trained and constructed transformers present the same patterns, further supporting the validity of our algo- rithm. Finally, theoretically analyze the predictor implemented by this construction (Sec. 4.3) showing that, in certain cases, it asymptotically converges to the maximum likelihood solution. Our findings provide insights into the mechanisms by which transformers perform model selection.
We defer to the appendix a discussion of related works (App. A), additional theoretical analyses, omitted proofs (App. B), extra experiments (Apps. D. F, C). and generalizations of our transformer construction (Apps. H. I, J).
2 (DISENTANGLED) TRANSFORMER MODELS
In the following we introduce the necessary background and notation about the models we use later. Transformers. The architecture of decoder-only transformers is built on two fundamental com- ponents, the attention mechanism and the multi-layer perceptron (MLP). Given a finite alphabet S, transformers map an input sequence = $1:7 (s1,..., ST) E S to a sequence of vectors 2 = (ZI....,,T) 2T) where z; E Rdd Each element of the input sequence s; is first encoded using its corresponding one-hot vector, cx, E 0, 1 These one-hot representations are then mapped to d-dimensional vectors via an embedding matrix E Rdx|S| To incorporate positional information, positional embedding matrix F E R is added. With a slight abuse of notation, let e; denote the i-th element of the canonical basis of IR" such that each input element is mapped to a vector r; R via Vi = Ee,, + Fe.. The information of the different tokens is then mixed by the causal self-attention heads: denoting the key, query and value matrices K, Q E Rdxdox V E I Rddd and given an input dxx one gets
Autn(h; Q, K) = A(h; K)hT, with A(h; K) = Softmax (Mi T QKTN);a),
exp (es/a) is applied row-wise and a > is a temperature parameter. where Softmax a); E, (mgfaa)
In the following, we call A = QK E Rddd the attention matrix, A E RTx the attention, and Attn Rdx* + R the attention layer. The causality of the self-attention is enforced by a mask M, to prevent the model from attending to future tokens, i.e. M(A);s = Aij if 2 j. -oo otherwise. For model with I. layers and (Hi [H])teLL] attentions heads per layer, we denote by Q(I,)), K(l,),, (((.)) the attention parameters for the i-th head in the I-th layer, w/"), wr Rdxd the parameters of the MLP at layer I, and Wo E IR (i the parameters of the output linear layer. Then, with (0) = (II,,..,44 IT) E RaxT as computed above, the decoder transformer (SIIT) can be written for = I..... J.....,, L. as
Hi Atm(h QU,,), K(U - + w (W/(TT00)
1 = (1-1)
where the output is given by Wol(L) RISITT
Disentangled Transformers. To improve the interpretability of the operations implemented by the models, Friedman et al. (2023) propose a transformer architecture in which each layer's output is concatenated, rather than added, to its input. This construction makes the residual stream explic- itly disentangled, but increases the embedding dimension (constant for standard transformers) with depth. Additionally, in such disentangled transformers the MLP layers are removed, the attention heads are parameterized by a single matrix À = QKT E Id xdz and the value matrices are ab- sorbed into the output layer Wo. Both the token and positional embedding are one-hot encoding, ie. E and F are identity matrices, and we encode the input 8; as [es,, 63] via concatenation rather than addition. Altogether, the disentangled transformer T(s1.) is formalized for 1 = 1,... L as (1,A)) Aum(h A(.))) for h = I...., Hi, and h (-1) (0.1) h A.
where the outputs is (t)( Due to the concatenation, the embedding dimensions grows layers as di = (1 + H,) - d- with dg = |S| + T. Importantly, Nichani et al. (2024) demonstrate that disentangled transformers are equivalent standard transformers using only attention layers.
3

---

Published as a conference paper at ICLR 2025
3 MARKOV CHAINS AND CAUSAL STRUCTURE SELECTION
To address the limitations of existing synthetic settings based on Markov chains and better capture the complexity of natural language, we propose novel framework. In this framework, the model must learn to select the correct causal structure in-context in order to solve the task and generate the input sequence. In the following, describe this task in detail and outline its solution.
Interleaved Markov Chains. The framework consists of sequences of length T on a finite alphabet of tokens S, generated by K distinct sources. Let 21 (UI..... UK) be the set of sources and K. = (A1,-- - k/}} set of positive integers; each source U, consists of ky interleaved and identical irreducible aperiodic Markov chains (Batu et al., 2004; Minot & Lu, 2014). All the sources are defined by the same transition matrix P pisix1SI, where P is the set of row-stochastic matri-
ces. This model is equivalent to a time- sime-homogeneous Markov chain (x(") )c20 of order ky. whose transition probabilities depend only single state k, steps back:
P(X;- - Xt-1 St-1, st- X1 =1) = P(X; - XI- - sc Sl-ky P
Here, we call kj E K. the lag parameter, as defined by Berchtold & Raftery (2002), where K S [1,4] is the set of possible lags. The lag, represented by the edges in Fig. 1, encodes the causal structure by explicitly representing the causal relationship between the variables the Markov chain.
Data. Given p and K, lag is uniformly sampled from K. for each sequence. Denoting the maximum lag by h max(K)), the first l elements of each sequence sampled from the stationary distribution क of p*, ensuring a constant number of independent variables for all sources. The
likelihood of sequence of lag k PP XT k) II x (X,) i PP(,, X)-++1).
Task. In this setting, the task is to predict the next state 81+1 given an imput sequence $1:7 generated from one of the sources, sampled at random. However, the identity of the source, and therefore the lag, is unknown. This task amounts to solving the following minimization problem:
S inf E k~-UnifflL...... DKL (IP(x++1 ((-x41) || I(XI, XT))
(x (rll..
where DRr. is the Kullback-Leibhe divergence. Eq. (1) admits a closed form solution which is the Bayesian model average (BMA), defined as the average of the transition probabilities for each lag. weighted by their posterior probabilities:
PP &)P(K)
P(X++1 X1:r) = Σ (xx(X Xx-++1) with we (XI:)) ΣNEK P(X1: B)P(K)
kEK
Asymptotically, the posterior distribution concentrates around the maximum likelihood (ML) esti- mate (Rousseau & Mengersen, 2011). Let k be the lag that maximizes the likelihood for a sequence (sI..... .sr) i.e., K* arg max P(X, probability wa converges to for k² and to 0 for the other lags, i.e., + [ = k=] where 1l is the indicator function. Then, BMA reduces to selecting the lag with the highest likelihood:
$1...... Xy sT k)). As T + oo, the posterior
Q(X+41 [1..................... (x)
I X Xx,++1).
REK
It is important to note that an interleaved Markov chain of lag k is mathematically equivalent to k-th order Markov chain with specific transition structure. Thus, given set of orders A. and sequence generated according to one such order, one could theoretically solve the task by learning in-context the corresponding (A: + 1 I)-gram transition probabilities (Nichani et al., 2024; Edelman et al., 2024). However, such an approach fails to leverage the low-dimensional structure of the problem, resulting in a suboptimal sample complexity of O([S!++1).
4 How CAN TRANSFORMER Do IN-CONTEXT SELECTION?
We now want to understand which algorithm transformers learn during training. We focus on disen- tangled transformers as defined in Sec. 2, which allow for more interpretable analysis of the model internal computations. The following proposition, which is the main result of the paper, shows how disentangled transformer can implement predictor solve the in-context selection task.
4

---

Published as a conference paper at ICLR 2025
Proposition 1. Let K. be contiguous subset of integers, i.e, K 1,, K |K] (nd k max For any T k there exists a three-layer disentangled transformer T with K heads in the second layer such that, defining Pi,e Eiex P. for
X X,
exp (T-K) 2xi=++1 Pi,k)
T(XIT)r Σ (((X)))((XHH1 Xz-k+1) with (xx()1:))= AEK
Σ exp mek (T- 2izk+1 Pi.m
(3)
The predictor implemented by transformers in Eq. (3) resembles BMA but differs in how it aggre- gates past information. Instead of using the posterior of each model as in BMA, our method employs weights proportional to the exponential of the average of normalized transition probabilities Pr. We analyze this predictor in Sec. 4.3 and discuss its convergence to ML. Proposition 1 illustrates how, for large B, transformers can implement selective induction heads, a mechanism that adapts to the input sequence by copying the token correspondent to the lag that maximizes the average normal- ized transition probabilities. The proof of Proposition I, involves an explicit construction for the weights of the disentangled transformer implementing the solution in Eq. (3) (an alternative third layer in App. H). Notably, this construction produces attention maps similar to those in standard transformers (Fig. 1), suggesting algorithm aligns with trained transformer implementations.
4.1 PROOF OF PROPOSITION 1: CONSTRUCTION FOR CONTIGUOUS LAGS
To aid intuition, we use a running example with visual illustrations for T = 10, K. = (1,2,3). We recall that each input element is encoded as (0) [e,,, es] E 1)
First layer: extraction of transition probabilities. The first attention matrix, A consists of two blocks: the first block operates on the semantic component of the input tokens, learning the transpose of the logarithm of the transition matrix. The second block (1) learns the causal relationships induced by each possible lag S + s; for k E K::
A
A =
log P
4
0
A
-2
+A
-1
1
4
-1
-
-
-4
-
-
-
A0
-
+1
+2
3
-
1
-
+> if i-jeK
44
44
43
-1
-1
-
4 = - if i-jØK.
+2
-1
3
-1
We can compute the first layer's attention as: (ea, e,]] A [ ez] = (log P)e,,, + Asign( (1) (A()
and applying the softmax (a 1): A (1 (0) A0) Erek allng Meneeen +x Eree (l0 ")amie--*
Pri wids for i 1,
For x + oc (in practice, for À large enough) and denoting Pi,K == Erek.rei Pai-ii
(0) (0: A')), {{ pi.i-j
if ifi-jEK
lim A A
1
ifi=j=1
0
elsewhere
Therefore, the output at index i after the first layer corresponds to a weighted average of the past tokens h for & E I. where the weights are given by the normalized probabilities p..:: R,)))
(0) for > and ( for 1. With the input vectors (0) being the concatenation of the one-hot encoding of the state and position [es,, 6i). the first [S] entries Pi for i >>>>>>>1 1 and , Esy- The remaining entries, due to the non-overlapping positional encodings, directly copy the normalized transition probabilities for the transition 8i-k - 8; into the |S| (i k)-th element of fi(1), To build intuition, refer to the example in Eq. (12) where the colors highlight transition probabilities of the same lag:
(0) (0) ;()) Exe Ps
of (1) correspond to 8; EKEK,A< Pi,ks,,
8s,2
R43
Jun
0
PaA
P4.
65.
A
Pass
Paz
Ps.
Pa1
Pb.
P.11
11
0
0
0
0
0
0
0
0
0
0
7,00)
0
B
Pa1
0
0
0
P5.1
P6.
0
0
66.2
01
0
0
11
01
(4)
Thoul
5

---

Published as a conference paper at ICLR 2025
The operation of the first layer is now explicit: for each token ( it extracts the normalized transi- tion probabilities Pi,k for each possible lag and stores them in the element h() TT The resulting vector is subsequently concatenated to the residual stream to be fed to the second layer.
Second layer: aggregation of transition probabilities. To predict the next token, the model needs to determine which lag generated the sequence based on the past transitions. This selection requires aggregating the normalized transition probabilities from the past, and storing them in the embedding of the current token. However, since consecutive tokens store transition probabilities in overlapping positions, the attention needs to learn convex combination of tokens that avoid mixing information from different transitions while maximizing the number of P stored (to not discard useful informa- tion). For instance, when aggregating the past for the token at 10 in Eq. (12), summing Agg) and Ri would mix P10,2 and 19.1 together. This mixing can be avoided, for example, by only select- ing tokens every 3 steps (1 R(", ) copying transitions without blending information. More generally, the attention (2.1)) should attend to every K-th token from the current one, which is equivalent to having non-zero entries along the diagonals at positions nK for n E N and nK < T. This structure can be enforced by constructing the attention matrix ( with single non-zero block operating on the tokens' positional encoding, as follows:
0 0
0
0
0
1
0
1
(2.1)
0
where the first & rows and columns are empty because the first elements of the sequence are sam- pled independently from the stationary distribution and therefore not informative. This construction resolves the issue of overlapping transitions, but copying only a subset of tokens implies losing in- formation from the excluded hú. Introducing additional attention heads _(2.2),.... A(. with the same form as (2211) above overcomes this limitation. The resulting attentions (4222) ((2.)) still follow a diagonal structure as ( to avoid overlapping transitions, but they are shifted to copy different tokens. For the given example, we can design (2,2) as in Eq. (5) to attend to (11)
(2.HH))
and (1)), and similarly, construct (2233 for ii and ):
A
A
A
(2222)
1
4
(22,3)
(5)
-
-
-
3
-
4X
3
3
-
1
-
+1
3
-
Each head has a diagonal structure with non-zero entries along the diagonals at position nK + h - for n 2 0 and h E [l...., H2}. the attention matrices can be formalized as:
+1, ifi > i and (i j) mod K = h
4 (2 = -1, otherwise,
where the condition i > j >> k ensures that all entries in the first k rows and columns are set to -X and imposes lower triangular structure due to causal masking. The modulo operation instead assigns each diagonal multiple of K to +λ (allowing attention) and the remaining diag- onals to -X (masking attention), while determines the shift of the the first positive diagonal to ensure the heads do not overlap. The output of each head in the second layer is given by [[e,,, ei],
(A,3)))))) Applying softmax and in the limit x + oo, the rows of the attention become uniform for positive entries and zero otherwise:
ey], R = A!2A))
A(2.A) = (A 1)g- = 1 [A@- 2 E m-1 1 [1(2)) = =x)=-'
(2.h)
The output i. (2.2) of each head is then concatenated into the residual stream. Fig. 2 shows the output for the 10th token for each attention head (2.11) (1(22) and (2233) to visualize the mechanism

---

Published as a conference paper at ICLR 2025
1%
PIA
.00.0
2162
Figure 2: Visualization of the mechanism of the second attention layer in construction. The trix represents the input of the second layer (2) whereas the single vectors the output for the 10th token (,(2.1) ((2.2) (2(2.3) Each of the three attention heads (arrows of different colors) copies non-overlapping transition probabilities at distance 3 from each other from the past. By doing this, the output of the second layer for the current token (10) contains all for each lag without loss of information.
of the second layer. The arrows of different colors represent how each head aggregates transition probabilities by attending to non-overlapping past tokens and averaging them with uniform weights.
When concatenating the output the different heads n, [h. ()) hg (1) R(2.1) ...., .,,,,,2 A,2.Ha)] we can how the 10th token stores the transition probabilities of its entire past for each lag.
Third layer: average of transition probabilities and lag selection. The third layer sums the normalized transition probabilities from the second layer embeddings and uses the result infer the correct lag. This mechanism is implemented through the combination of multiple blocks within the third attention matrix, A (3), which is structured as follows:
A(3)
A(3)
B 3
+1 fiij+lek -1 if if-ji+
B(3) Gg = 3 5 0, otherwise
(+1, (1, if (i- mod = K
The matrix A3) acts the positional embedding of the input, similarly the matrix A in the first layer. The difference is that the position of the diagonals is now shifted by one. This shift ensures that the only non-zero entries after softmax are the
on the diagonals corresponding to - for A K. The matrix (3) is instead responsible for the sum of the normalized transitions. Each block operates on the output corresponding head in the second layer. To understand how, consider the following tokens in output of the first head in the second layer,
Va
0
Ac
P4,
Pe.
f4,11
Pr.a
91o,3
Pro.2
P1o.1
k/2.4)
Fing
Ao
Ne.3
Av,
Av.1
84.3
V/
D
D
0
8s
0
Ps.a
16.
Ps.
Ps.1
abilities and m (2 (2,4)) where we define D
E R² as the block of i. (2.h) which contains the normalized transition prob- E R contains a copy of the second attention. With the structure of A3),
can see how (3) acts on these two blocks: A (2)T A(3) (2) Σh=1P
(2.A)T ((3) m. (2.h) + (3) eg- This operation sums the transition probabilities such that the entry corresponding to the lag-k tran-
sition for the next token contains the sum of the transitions with the same lag: !!22T A 1-

---

Published as a conference paper at ICLR 2025
Ejs: Pjs. To illustrate this process, consider the following:
Pia
(2.1)T g(3.1
(2.1)
B
Phaa
Ph.
Ph.3
The additional blocks containing B on the outputs of the other heads, performing the same operation by summing the transitions of the same lag stored in the respective outputs. Considering all heads and only the non-zero entries after softmax, occurring at j = i-ik + due to A we get
K
K
This
A. (2) A |
P. (2
B(
(2.M) m b+1
Pk PkthtnK,k + A,
Th
where Th L h Taking the softmax (a 1) and the limit À + co, then considering the last T-k and absorbing K into the temperature 8, we recover
token T (for large 7) where Th, + 1 ~ weights in Eq. (3):
K
A (3) (Mi (((-K+11) = exD (T-
imkil
Pi.k
Σ
exp
(T-K)
1=2+1
E B..r): for k E K.K.
(6)
4.2 SELECTIVE INDUCTION HEAD AND NEXT TOKEN PREDICTION
Selective induction head. Eq. (16) shows how the last attention layer computes weighted average of the tokens at a distance & from the next one, with weights proportional to the average of the lag k. In practice, trained models often learn large values of 3. Thus, we consider the limit 3 + co, where the softmax converges to the hardmax:
normalized transition probabilities
0 - k*] with k- = argmax, Ejer Pj.
1+++=
Here, the transformer selects the causal structure (i.e., the lag) corresponding to the largest Ejes Pj.k. Given the current token i, the third layer then copics the token from the position i- -**+11
i.e., n,()) Σj= [i 1 k*) ()) 1 + After concatenation to the residual stream, the tokens of the following form:
h
(3) Cocs e, (1) (1(211) H2. 1- 1
(())
8 1 +1
(2.H2) -
Output layer: next token prediction. Finally, the output layer Wo E S contains all zero blocks, except for the one acting the semantics of the token copied by the third attention. This block learns the transition matrix p to predict the transition probabilities to the next token via
Wo = (0sxs1 OSxT 0sxdo Osx2do
Osx2d,
P
0sxt
Osxdp
0gx2do .... 0sx2do)
i.. (oh = p* 6 = P 4, This layer shows how transformers can learn selective induction head; a mechanism that adapts to the input sequence by copying the token corresponding the argmax of some quantity extracted by the previous layers and stored in the embeddings.
4.3 EQUIVALENCE WITH MAXIMUM LIKELIHOOD
The disentangled transformer we propose does not rely likelihood. Due to the normalization applied by the softmax function, the model computes the sum normalized probabilities p. For the inference to be accurate, the cumulative sum corresponding the correct lag must exceed those of any other lag. This fact is formalized in terms expected values in the following claim: Claim 1. Let K. be subset of integers and X, an interleaved Markov chain of lag K, then, for E KC and > k.
XL P'X;
XTT P
X.T
E
SE
ΣIEK PLP,
ΣieK
XTIPP*X
8

---

Published as a conference paper at ICLR 2025
R= (1, 2,
X2 = (1, 3, 4}
R= (1, (1,2,3, (1,,,, 4,
R = {1,3)
10a
KL. Divergence
70ni 311
Frote 3H
Foun 2tt
10-1
Periis I11
Perbs 612
Trasee 211
Tracett 317
Proesir 301
Tonut:: 121
10
Prcestr
Tonutts
10
10-1
50 100
Sequence Length
Sequence Length
Sequence Length
Sequence Length
Figure 3: Performance of our constructed transformers, trained transformers, and theoretical estimator (BMA, ML). First plot: lags 1.2.3. Second: the model solve the task with non-contiguous lags. Third: the model is effective with additional lags. Fourth: one head is enough for two lags.
While specific cases (e.g., two lags, no normalization, or independent lags) are proven in App. B, we leave the complete proof of Claim for future work. However, we provide empirical validation of this claim in App. F. Due to ergodicity, the average + - Pe,r converges to its expected value, and for large enough 7' it is higher for the correct lag. Therefore applying the exponential and scaling the temperature / leads to the same result as MLE in the asymptotic limit, as shown in Fig. 4c.
4.4 GENERALIZATIONS AND SPECIAL CASES
Single-head transformers: the construction above allows the model to store all the past transition probabilities in the embedding of the current token by scaling the number of heads with the total number of lags K. Since all heads perform the same operation, reducing the number of heads im- plements an equivalent algorithm but with worse sample complexity compared to MI., as some past transitions are discarded. Thus, a 3-layer single-head transformer still solves the task by implement- ing the algorithm in Proposition 1, but it only uses T-A-1 + 1 samples to estimate the correct lag. Non-contiguous lags: in App. we provide examples of constructions handle non-contiguous lags, where the core approach remains similar to that in Sec. 4.1. Depending on the specific case, the number of heads needed ranges between the number of lags and k min(AC) + 1.
K
Two lags: by the construction in Sec. 1. handling two contiguous lags requires two attention heads in the second layer for optimal sample complexity. However, we provide in App. J an alternative construction for the third layer which enables a single-head model to match the performance of the two-head model for any two lags, whether contiguous or not (see empirical results in Fig.3).
EXPERIMENTS AND DISCUSSION
We conduct a series of experiments to empirically validate our construction and determine whether transformers trained via gradient descent learns it. Setup. We train 3-layer disentangled transform- ers (7) and 3-layer standard transformers (T) with learned positional and semantic embedding both with a /dok using cross-entropy loss. At each step, we generate fresh batch (size 256) of sequences (length 128) via Alg. 1, and train using Adam optimizer with fixed learning rate 0.001 and no weight decay. For the standard transformer embedding size we tested 128, 64 and dok = 32. For the constructions, we fix / = 100 and x = 500. We report DKr. between the true and predicted next-token distribution along the sequence. We generate different tasks with alphabet size |S] (no differences are observed for other sizes) varying the number and values of lags: K. (,) (our example) and IC 21, 2, 3, 53 for the case of contiguous lags (optimal number of heads and 5 according to Proposition 1), K. = 31,3, 4) to show non-contiguous lags (4 heads needed, see App. 1), and K = (1 for the special case of two lags.
Main results. We observe in Fig. how the construction with optimal number of heads, indicated as T constr., matches the performance of the ML. Morcover, both the disentangled transformers trained ( train) and the standard one (T train) match the performance of the theoretical construction. Interestingly, we observe that when the number of heads is fixed to 1 the trained transformers can find solutions which perform better than the construction: this indicates the existence of more efficient, yet difficult to interpret, ways of aggregating the transition probabilities. and that gradient descent can find them. Finally, we illustrate how for the case of two lags (right in Fig. 3) construction

---

Published as a conference paper at ICLR 2025
(b)
Trited trandormer
X² (3,5,7>
Kt. Diveerggen
20-1
20-3
Trained disentangled transormer
10
Sequence Length
(c)
Dec K= (1,
disentangled
p-22
2-180
10-1
L0-2
Kl. 10-1
Ataaitn Layne
50
100
Sequence Length
Figure 4: Attention maps trained vs constructed transformer. Fig. 4a reports the heatmaps of the attention maps of the trained standard and disentangled transformer and our construction for lags 3,5,7 and Fig. 4b. The first and second layers display remarkable similarity and the second layers show similar diagonal structure. Fig. 4c shows how estimator Eq. (3) matches ML for large B.
with single head (detailed in App J) attains the optimal sample complexity, while the construction in App. would require 2 heads. Moreover, in this case, it appears that the trained transformers can obtain performances closer to the BMA rather than the ML for small sequence lengths.
Trained vs constructed attention maps. In Fig. 4a report the attention maps for the three layers for {3,5, 5, 7] for trained standard transformer (top), trained disentangled transformer (middle) and our construction (bottom, construction in App. 1).The attention maps of the first and third layers are nearly identical between the trained and theoretical models, with these layers functioning precisely as expected from the theoretical construction. Notably, the attention entries of the first layer are proportional to log P* even when the model is trained from scratch, for both the disentangled and standard cases. For the second layer (aggregation), the trained transformers con- verge to a slightly different structure, likely because aggregation is a combinatorial problem with multiple valid implementations. Nevertheless, clear diagonal pattern emerges, closely resembling our construction. Furthermore, as demonstrated in Fig. 4b, all models achieve comparable perfor- mance on the task. These findings strongly suggest that the trained transformers find a solution that aligns closely with our construction. Remarkably, we also show that standard transformers trained with learned positional and semantic embeddings and attention parameterized by Q, K, V produce attention maps in agreement with our construction. This provides compelling evidence that our construction is not merely a byproduct of the disentangled transformer's 's architecture but it is also implemented by standard transformers. Interestingly, even with embedding dimension (64) smaller than the sequence length (128), standard transformers efficiently store and utilize past transitions while matching optimal performance.
CONCLUSIONS
We introduced a novel synthetic task based on interleaved Markov chains to study how attention- only transformers perform in-context causal structure selection. Our findings demonstrated that 3-layer transformer can solve this task with near-optimal sample complexity, effectively showcasing the emergence of selective induction heads, attention circuits that aggregate past information and select the correct causal structure. Moreover, we provided a fully interpretable construction of disentangied transformer implementing these circuits to solve the task, and empirically verified that both disentangled and standard transformers trained with Adam closely align with this construction. Finally, we theoretically analyze the algorithm implemented by this construction showing that, in certain cases, it asymptotically converges to ML.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work supported by the Swiss National Science Foundation (grant number 212111) and partially funded by an unrestricted gift from Google.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neu- ral Information Processing Systems, 2023. URL httpt:ttt///t...t. https://openreviit /ttp///ppppr/iieeee///ee/ee wwtpenrwwwwwweeeeeeeeeeeeeeeeeee www.foorrrmmmme forum? d=
L, mrinnii 9...
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:221 15661, 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
Tugkan Batu, Sudipto Guha, and Sampath Kannan. Inferring mixtures of markov chains. In Learn- ing Theory: 17th Annual Conference on Learning Theory, COLT 2004, Banff, Canada, July 1-4, 2004. Proceedings 17, pp. 186-199. Springer, 2004.
André Berchtold and Adrian Raftery. The mixture transition distribution model for high-order markov chains and non-gaussian time series. Statistical Science, (7(3):2283366 2002.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transforer: A memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Benjamin L Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution statistical induction heads: In-context learning markov chains. arXiv preprint arXiv:2402. 1004, 2024.
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. In Thirty- seventh Conference Neural Information Processing Systems, 2023. URL https:// openreview net/forum htte:/oormmmiddePelee wid=eeeWxkNeeP
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 55:30583-30598 2022.
Atticus Geiger, Chris Potts, and Thomas Icard. Causal abstraction for faithful model interpretation. arXiv preprini arXiv:2301.047000 2023.
Jiachen Hu, Qinghua Liu, and Chi Jin. On limitation of transformer for learning hmms. arXiv preprint arXiv:2406.04000 2024.
Muhammed Emrullah Ildiz, Yixiao HUANG, Yingcong Li, Ankit Singh Rawat, and Samet Oy- mak. From self-attention to markov models: Unveiling the dynamics of generative trans- formers. In Forty-first International Conference on Machine Learning. 2024. URL https: / wepprreview.nee/oooroeoeeeeeeeee 4mp LUb.
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.
Ariana Minot and Yue M Lu. Separation of interleaved markov chains. 2014 48th Asilomar Conference Signals, Systems and Computers, pp. 1757-1761. IEBE, 2014.

---

Published as a conference paper at ICLR 2025
Aaron Mueller, Jannik Brinkmann, Millicent Li, Samuel Marks, Koyena Pal, Nikhil Prakash, Can Rager, Aruna Sankaranarayanan, Arnab Sen Sharma, Jiuding Sun, et al. The quest for the right mediator: A history, survey, and theoretical grounding of causal interpretability. arXiv preprint arXiv:2408.01416 2024.
Timothy Nguyen. Understanding transformers via arXiv:2407.12054 2024.
n-gram statistics.
arXiv preprint
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference Machine Learning, 2024. URL
bitts://cpeerrett
net/ Forum?id= JNM4 imlHZv.
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. 11895, 2022.
Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vard- han Makkuva. Transformers on markov data: Constant depth suffices. arXiv preprint arXiv:2407.17680 2024.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
Judith Rousseau and Kerrie Mengersen. Asymptotic behaviour of the posterior distribution in over- fitted mixture models. Journal of the Royal Statistical Society Series B: Statistical Methodology, 73(5):689-710, 2011.
Michael Eli Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How trans- formers perform in-context autoregressive learning In Forty-first International Conference on Machine Learning. 2024. URL https::/p/// /
Naomi Saphra and Sarah Wiegreffe. Mechanistic? arXiv preprint arXiv:2410.09087 2024.
Anej Svete and Ryan Cotterell. Transformers can represent 7--rrm language models. arXiv preprint arXiv:2404. 14994, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL atttps://proceedding wwwrips..c...oo hatpaadddarr..arrr....oocappprrr wttp://eeeeeeeeeeeaepepeeeeeeeee files/paper/20117 File/3f5ee2432 2eee2443547de9991 44t.444455....100000001444444aa3
pdf.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023a.
Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309. 05858, 2023b.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. arXiv preprint (eXiv:2/11.02000) 2021.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models, arXiv preprint arXiv:2311.00871, 2023.
12

---

Published as a conference paper at ICLR 2025
Organization of the Appendix. The Appendix is organized as follows. In App. A, we present detailed review of related work on in-context learning and mechanistic interpretability. App. B ex- tends the statistical analysis of the estimator implemented by the transformer in Prop. and includes omitted proofs. App. C reports additional experiments and discussions about using more than 3 layers and varying the number of heads in the second layer. Additional attention maps for differ- ent tasks and both disentangled and trained transformers as well as the construction are provided in App. D. App. includes several experiments validate Claim 1. App. details the algorithm used to generate the interleaved Markov chains. A complete construction for the contiguous case, along with detailed intuitions, is given in App. G. App. H discusses an alternative third layer construc- tion using positional embedding. The construction for non-contiguous lags is presented in App. I. Finally, App. explains the single-head construction for the case of two lags.
A RELATED WORK
Following the initial empirical observations of the emergent in-context learning capabilities of trans- formers (Brown, 2020), several works have attempted to understand this phenomenon. Xie et al. (2021) sought formulate in-context learning as Bayesian inference, while Garg et al. (2022) stud- ied the ability of transformers to learn simple functions, such as linear models or multilayer percep- trons, context. A subsequent line of work (Akyürek al., 2022; Bai et al., 2024; Von Oswald al., 2023a;b; Raventós et al., 2024) shows that transformer layers might implement gradient descent to solve in-context linear regression. Ahn et al. (2023) extend this idea to higher-order algorithms. Im- portantly, Olsson et al. (2022) postulate that in-context learning is tied to the emergence of induction heads. Bietti et al. (2024), subsequently extended this idea, showing the development of induction heads to learn bigrams in-context and showcasing connection with associative memories. Our work is inspired by foundational efforts in causal interpretability and mechamistic understanding of neural networks, particularly transformers (Olsson et al., 2022). These include studies causal mediation analysis (Mueller et al., 2024), causal abstraction as a theoretical framework (Geiger al., 2023), and the broader scope of mechanistic interpretability Wiegreffe, 2024). More closely related to our work is the literature analyzing transformers through the lens of Markov chains. In particular, Nichani et al. (2024) shows how transformers trained on sequences generated by Markov chains on graph learn simple circuits to capture the underlying causal structure and implement the Bayes-optimal solution by estimating transition probabilities in context. Similarly, Edelman et al. (2024) illustrate the formation of statistical induction heads that accurately compute posterior probabilities based on bigram statistics. Makkuva al. (2024) used Markov chains to study the loss landscape of transformers, while Rajaraman et al. (2024) shows that a constant depth is sufficient to learn k-th order Markov chains. Along this line, Svete Cot- terell (2024) demonstrates that transformers with hard sparse attention can exactly represent any n- n-gram model. Extending this framework to Hidden Markov Models, Hu et al. (2024) highlights the limitations of transformers in learning such models compared to RNNs.
language models (Saphra &
B STATISTICAL ANALYSIS OF OUR PREDICTOR
For our estimator to select the correct lag, the following inequalities must hold for lag k and a sequence of length? generated accordingly:
T
Σ PI.k Σ Pj.r V h: and; E K.
These results enable us to recover the MI.E estimator in the high-temperature limit and approximate the BMA at finite temperatures. Assuming the process is ergodic, and by taking the limit of the inequality above, we require the following condition:
X',,PXX
K.T.P"X.
E
S.
for r E 1C and max((C),
ΣieK PL
ΣiEK X P'X.
as formalized in Claim 1.
We leave the complete proof of this result as future work, but we have fully validated it empirically in Section F. We provide here the proofs of Claim 1 for three specific cases.
13

---

Published as a conference paper at ICLR 2025
Two-lag case. In the case of two lags, can show the following general result for any two distributions, P and Q, over 1,...,k k for lo 2
Lemma 1.
k P(OQ()
P(1)
P(1) QQ1)
P(i) + Q(1)
Proof. We first show that
k Q(1)=
(P(1) - (()
k (PP() Q(1) 0.
Σ [=1
1=1 +()+Q())
P(t) Q(1)
2=1
Then, by using Cauchy-Schwer inequality, we obtain:
P()Q(() +(1++(11)
P())
Q(1)² P(i) + Q(1)
P(1) P(i+ +
P(1+ + Q(1)
The result in the two-lag case follows directly, Let fe denote the distribution of the lag-k interleaved process Xx. (i.e, p(s;, 8j, 8k) P(x; 8t, Xj Sj. Xx sk)) For any lag {r} we have
X P*X;
Psi P Pol-risc
Σ
ki 8i)
P + X
Pai + PS
kid
Pol Pei
((8i-k,8i-r)
Psi P Pai-rine
By applying Lemma 1, we directly obtain
XT TTP.XX
X PM X; X Px
E
E
XX X P.
P X,
which proves Claim in the case of two lags.
Independent lags. In the case where all lags in K. are such that (Xi-i)tek are independent, can prove Claim 1. Indeed, in this case, the distribution of the observed lags can be factorized as (((Ss-))EE)) [[tex (() Thus we have
XT P'X, -X X P'X. [
P Painkio (Poinie Pa PaI_Rios)
Σ oulee-t)iek
(((Ss-))EEX)
ΣieK P.X,
Σiek Penniis
Psr nis, n.,,(Pa rit - Pas m.s.))
Σier. Poi
Then, by observing that a(a b) + b(b - a)= = ((a--)- b)22 2 (a-
the result follows from:
P.. kid (P - -P.
2 Σ
u(s:-))u(s(-))
ΣiEK P. PoI_iee
Pai e. (Po: - Pai mon)) + (P u:s; - -Pa
Σ M(8)-))((81-))
Eiex Poi
81-
(Prinrss P Eiex Pss use
Σ ((siz))((i-4 S(-P,81-1
20
We observe that similar techniques can be applied to prove Claim in the case of a symmetric Markov kernel P*.
14

---

Published as a conference paper at ICLR 2025
No normalization case. Due various reasons, including the normalization of the softmax in the attention layer, our estimator relies on a score computed by aggregating the normalized probabilities Pi.a- If we were to use the unnormalized probabilities, we could rely on the following result, which simplifies Claim by excluding the normalization step.
Lemma 2. be a subset of integers and X, a stationary interleaved Markov chain of lag E K, then
B|X_,, P X1] E|
for r T K and i 2 max(K))
(8)
Proof.
B B B PxX
Σ
k;
84 ki
Σ
1( 8i PoLKke ParieePaiii Pai-ris
11-11
1
Σ
((
Σ Si-b) sc-r.8 Ph
knd
8;
Si-Py8i=b
Ni-riSi-R
$ Σ
Σ ((8--1;
Σ ((- pr
kid
Si-h
8i-m 8i-8
where the inequality follows from the Gauchy-Schward inequality. Assuming that u($;-r, si is coupling of the stationary measure then have:
B[X__, P
x(s:k))P2
Σ x (s
p2
8t.8g
1E P X.].
remains to prove that (18--n, 8i is coupling of the stationary measure æ.
First, let's assume that and & are such that X;-r and X;-k are independent. In this case (1 88-x) ((((-))((s;==) x) and we have both that Σ.; (((s-r:8(-x)) = p(s;-)) and E (((8)-188(-1) = (())-r).
Alternatively, if r and & are such that X, and X come from the same Markov Chain with r k. We have thus that X,-r p and X,-t|X,-r ((-)-r:3((k)) = (8;-r)ss P' Since P is a stochastic matrix, summing over 8;-k gives Ee.r p(si-r, 3i-k) = u(si-r) = ((s_r).. Finally, by definition of the stationary distribution m summing over $i-r gives sier p(st-r, st-k) x()-k)).
sT,Pl for some 1 2 and
15

---

Published as a conference paper at ICLR 2025
C SCALING HEADS AND LAYERS
In this section, we investigate how varying the number of heads the second layer and the number of layers affects the model's performance. We train standard transformers with learned positional and semantic embeddings in the same setup as reported in Section 5. In Figure 5 (left) we consider the task given by K {1, 2. and first show the behavior of the model with layers and different combinations of heads [1. 11. [3, 1]. [1.3]. [3, 38,3"", the results show that transformers with 2 layers can't t solve the task. Second, we show that increasing the number of layers beyond does change the performance. In Figure 5 (right) instead we consider the task defined by K. {1, 2,,2,3, 4,5) and train transformers with fewer, equal to, or more than K heads. As predicted by our construction increasing the number of heads leads to performances that get closer to the maximum likelihood up to having the number heads equal to the number of lags in the set K. Beyond this point adding more heads does not improve performance, this is expected as ML is optimal. Figures 6a,6b,6c illustrates the attention maps for a 3-layer transformer with only 1, 2, heads respectively in the second layer, despite the task having lags. Remarkably, even with fewer than K heads, the layers remain consistent with our theoretical construction, displaying analogous patterns: the first layer extracts transition probabílities, the second aggregates them, and the third implements the selective head. However, in the case of fewer heads, the second layer appears to find an efficient way to superpose information-2 a mechanism we could not yet interpret. Understanding this behaviour in the second layer remains an open question for future work.
k= =
K =
1,, 4,
10
10
KL Divergence
KL Divergence
10-1
10-1
10-2
10
10-3
10-3
10
10
50
100
50
100
Sequence Length
Sequence Length
BMA
311
BMA
-1, 2, 51,
721, - (1.3..1
F3L 3. (1.66.
[2L
72L [3, (5.
73L I1.
73L 119L1,44111 4]
11,
Figure 5: (left) Scaling number of layers: we train standard transformers with learned position and semantic embeddings. Transformers with 2 layers can't solve the task for any combination heads. Transformers with more than 3 layers achieve the same performance as for 3 layers. (right) Scaling number of heads in the second layer: we train standard transformers with learned position and semantic embeddings increasing the number of heads in the second layer. As predicted by the construction increasing the number of layers leads performance closer the Maximum Likelihood.
With this notation we intend the following [#heads layer 1, 1,... #heads layer I.]
16

---

Published as a conference paper at ICLR 2025
Attention Layer
Attention Layer
Attention Layer order
75
Seepontoo 2rdee
20
Suquinee
(a) Attention maps for TL [1, 1, and A= (1,2,3,4,5},
Layer
Layer
Layer
Sqqunee Taba
Sequeenee Iodes
15
30
(b) Attention maps for 7 3L 21, and K.
Attention Layer
Attencion Layer
Attention Layer order
Seepueoo Todes
20
20
25
(c) Attention maps for T 3/ [1,3, 1] and 21,2,3,4,5)
Figure 6: Attention maps for different heads in the second layer and K. (1,2,3,4,5): observe how even with fewer heads the transformer learns layers which are consistent with the operations in our construction. particular the first layer still extracting the transition probabilities, the second is aggregating them and the third one implements the selective head.
17

---

Published as a conference paper at ICLR 2025
D ADDITIONAL ATTENTION MAPS PLOTS
As an additional confirmation for our construction, we report here comparison of the attention maps after softmax for the task introduced in Figure We compare, trained standard 3-layer attention- only transformer with learned positional encoding and one attention head per layer, trained disen- tangled transformer and our construction. The standard transformer was trained in the same setup already introduced in Section 5. In Figure 7 we train on data with KC = 1,2, we observe remarkable similarity between the attention maps of our construction and the trained transformer. This further confirms that the disentangled transformer is good proxy to study the residual stream and the flow of information inside the transformer construction is realistic and aligns with what transformers learn in practice by gradient descent. Moreover, In order to showcase the adaptivity in-context of the selective induction head depending on the true lag of the input sequence, in Figures 8, 9, 10 train lags AC f1,2,3) and test on sequences generated with each one the training lags, similarly in Figures 11, 12, 13 we train on lags K {1, 3) and test on each. As expected, the third layer adapts the input sequence selecting the correct lag and copying the correspondent token via the selective induction head.
more interpretable way. Moreover, it confirms that our
Trained standard transformer
Sequarce tatex
Trained disentangled transformer
Saquanoo Intes
B
Construction disentangled transformer
Sequente indes
5
Attention Layer
Attention Layor
Attontion Layer order
Figure 7: Attention maps KC = (1,2) and true lag A: 2.
18

---

Published as a conference paper at ICLR 2025
Trained standard transformer
Trained disenlangled transformer
Construction disentangled transtormer
Layer
Layer
Layer
Figure 8: Attention maps
(1,2,3) and true lag k 1.
Trained standard transformer
Trained disentangied transformer
Construction disentangled transformer
Layse
Layer
Luyer
Figure 9: Attention maps K (1,2,3) and true lag l: 2.
19

---

Published as a conference paper at ICLR 2025
Trained standard transformer
Trained disenlangled transformer
Construction disentangled transtormer
Layer
Layer
Layer
Figure 10: Attention maps K. (1, and true lag k
Trained standard transformer
Trained disentangied transformer
Construction disentangled transformer
Layse
Layer
Luyer
Figure 11: Attention maps KC (1,3,4) 4) and true lag kc 1.
20

---

Published as a conference paper at ICLR 2025
Trained standard transformer
Trained disentangled transformer
Construction disentangled transtormer
Layer
Layer
Layer
Figure 12: Attention maps K. (1,3, 4) and true lag k
Trained standard transformer
Trained disentangied transformer
Construction disentangled transformer
Layse
Layer
Luyer
Figure 13: Attention maps AC (1,3,4) 4) and true lag k
21

---

Published as a conference paper at ICLR 2025
E TASK DETAILS
In this section, we illustrate the algorithm we used to generate batches of new samples at each iteration.
Algorithm Generate Dataset of N Sequences from Interleaved Markov Chains
Require: N (sequences), T (length), S (state space), (set of lag values), P (transition matrix) Ensure: Dataset D of N sequences
D <
2: π - stationary distribution of P
to N do 3: for
4: k <- Uniform(AC)
D Randomly select lag for this sequence Initialize first state Initialize sequence
5:
Xo f- Sample from æ
6:
S [Xo]
7:
for = to T do
8:
if f < & then
9:
X₁ + Sample from π
b Sample from stationary distribution
else
11:
X, + Sample from P* [x
b Transition based on lag
12:
Append X, S
Add S to D return D
F EMPIRICAL VALIDATION OF CLAIM
To empirically validate Claim we first sample of 12 lags uniformly between 1 and 30; then sample 1000 different transition matrices and for each matrix and each lag 1000 sequences of length 1000 according the respective Interleaved Markov chain. For each lag and each of sequences we then compute the expectation in Claim by averaging the last transition in each sampled sequences. We then compute the following quantity:
P'X, DIEK P*X;
X'P,
max E ritk TEX
and report it in the histogram in Fig. 14. We can see that all values in the histogram are positive therefore confirming our claim. Similarly, the results in Fig. 15 report the quanity in the claim for each single lag. As per our claim, the expected normalized transition probabilitics of the true lag is always larger than the same quantity for any other lag. As further confirmation of the claim, in Fig. 17 and Fig. 88 we report the cumulative average of the normalized transition probabilities along the sequence for single sequence. We observe that even with few samples (small the cumulative average for the true order is always larger than the same quantity for the other lags.
22

---

Published as a conference paper at ICLR 2025
1250
1000
Frequency
750
500
250
0
0.020
0.025
0.030
0.035
E x__ P'X,
P
max E rEK:
Eier x,
P.X.
Figure 14: Difference of the expected Normalized Transition Probabilities for the true lag and the max- imum over all other lags for [S] = 10. The histogram shows how the quantity in Eq. 9 is always positive.
Thu lay
Trus lag:
True lug
Tran bay
Tour Iag
3.10
3000
1100
True 1ug
10
True 1sat
The lop
Troe lig
3.100
True lag 15
True Ing: 22
Drue Joz
Troe log:
2.15
3000
Figure 15: Expected Normalized Transition Probabilities for |S| 10: The sampled of lags is K. (1, 7, 10, 11, 13, 15, 16, 22, 26, 28), we sampled 10 different transition matrices and for each lag and each matrix sampled 1000 sequences of length 1000. The expected normalized transition probability always larger for the true lag.
23

---

Published as a conference paper at ICLR 2025
Thue
Troe le::
True Ing:
Tran
1110-
1000
Trar lag 10
Tene Tage
Thu
Tran
1.15
2.10
Tree Jug 15
Trme legt 22
The
Troe lag:
3.100
2.10
Figure 16: Expected Normalized Transition Probabilities for |S] 25: The sampled set lags is K. 21, 7, 9, 10, 11, 13, 15, 16, 22, 26, 28). we sampled 10 different transition matrices and for each lag and each matrix sampled 1000 sequences of length 1000. The expected normalized transition probability is always larger for the true lag.
Tiua Lag
Thei 2ag
The Tay 0
Tr 1e
Tix Fay 13
Tlar Day 15
Lay
Tue Lay
Sec B0b0
Drgi Far M0c7
Scde
Figure 17: Cumulative average of Normalized Transition Probabilities for |S] = 10: The sampled set of lags is K. 2,,2,7,9. 10, 11, 13, 15, 16, 22, 26, 28), report one sequence sampled according one the transition matrix. The cumulative average normalized transition probability quickly becomes larger for true lag.
24

---

Published as a conference paper at ICLR 2025
Thu Teg
The Eay
To See
Tixe Eag
Tux Lar
+143
Digilie
Figure 18: Cumulative average of Normalized Transition Probabilities for [S] = 25; The sampled set of lags is K (1, 7, 9, 10, 11, 13, 15, 16, 22, 26, 28), we report one sequence sampled according to one the transition matrix. The cumulative average of normalized transition probability quickly becomes larger for the lag.
25

---

Published as a conference paper at ICLR 2025
G COMPLETE CONSTRUCTION FOR CONTIGUOUS LAGS.
First layer: extraction of transition probabilities. The first attention matrix, A consists of two blocks: the first block operates on the semantic component of the imput tokens, learning the transpose of the logarithm of the transition matrix. The second block A(") learns the causal relationships induced by each possible lag s 7 for k E K::
log P
A
A(1)
A
{+} if i-jEK.
4
(-X if i-jęK.
4A +A
We can compute the first layer's attention as:
[ea,, ei] [ea, ez] (log P)ayie, + Asign (A!))
and applying the softmax:
(log Plagins Asigní ( (19)
flog P)oyice i (
A0) (0) (1 A
Er=1 (log Ploye, Asign( (1)
Erek
eloR P) Plon.es + + Erex elloR Plonees -
Poikin
For + oo (in practice, for large enough) and having denoted Pi == Erexieer Puiriie
for > 1,
Priiss if ifi-jE ifi=j=1- elsewhere
lim X
A
Therefore, the output index after the first layer corresponds to a weighted average of the past tokens hit for k E IC where the weights are given by the normalized probabilities pi,::
E 1 jE
K]
Paiics (0) if 1
Attn Attn(h (1:0):
Σ
rrek Paris
(10)
if
Σ
ifi> 1
kek,k<i
(11)
11
if==1
With the input vectors h (0) being the concatenation of the one-hot encoding of the state and position [es,s ei). the first [S] entries of i. correspond to ă; EBEK Pi,k's,- for >> 1 and & = Cx.- The remaining entries, due to the one-hot positional encoding, directly copy the normalized transition probabilities for the transition 8i-k into the |S] + (i k))th element of h("), To build intuition, we refer to the example in Eq. (12) where the colors highlight transition probabilities of the same lag:
AuS
0
Paz
Pu
(1)
P1
0
0
0
0
0
0
0
11
11
0
0
0
0
0
0
0
(1)))
0
0
0
0
Pa1
P...
Paa
PL
0
(12)
26.
Ph.2
Piia
Ps.1
P9.
PIos
Pio
Zina
Pin.
The operation of the first layer is now explicit: for each token (,)), extracts the normalized tran- sition probabilities Pi, for each possible lag and stores them in the element (1)) sulting vector is subsequently concatenated to the residual stream to be fed to the second layer (1) [[e,,, (((2).
S The re-
26

---

Published as a conference paper at ICLR 2025
Second layer, aggregation of transition probabilities: To predict the next token, the model needs to determine which lag is more likely to have generated the sequence based on the observed tran- sitions up the current token. In order to do so, need to aggregate the normalized transition probabilities of the entire past and copy them into the embedding of the current token. However, since consecutive tokens store transition prohabilities in overlapping positions, simply averaging token's ss history via uniform attention would cause information from different transitions to mix. For example, when aggregating the past for the token at i 10 in Eq. (11), summing hgg) and ((9) would mix P1o.2 and Po.1 together. To avoid this, the second attention needs to use the po- sitional encoding to learn convex combination of the past tokens, which minimizes the overlap and maximizes the number of transitions stored. For example, by selecting one token every 3 steps (1)) (())) R00 attention A(2) should attend every K-th token starting from the current one. This corresponds to a diagonal structure, with non-zero entries along the diagonals at positions nK for n E and nk < T. This structure can be enforced by constructing the attention matrix Á(2 with a single non-zero block that operates on the positional embeddings and takes the following form:
copy all the transitions without mixing information. More generally, the
Q
0
0
0
0
0
0
0
0
0
0
0
(
(
A(2.1)
Q
0
0
0
0
0
0
0
A
0
0
where the first k rows and columns are empty because the first elements of the sequence are sampled independently from the stationary distribution and therefore no transitions are present.
The current construction allows us to overcome the issue caused by overlapping transitions; however, only copying subset of tokens implies losing the information contained in the remaining å, which were excluded by the attention. To overcome this limitation, we can introduce additional attention heads, for total of K. Each one of them still presents diagonal structure in order not to overlapping transitions but now shifted such that each head copies different tokens. For the given example, we can construct ((2.2) as in Eq. (13) to attend hg))) and hgg which were discarded before, and similarly, the third head can be constructed with (2.3) to sum i and hg This mechanism corresponds to having each additional head with a diagonal structure and non-zero entries along the diagonals at position
nk + for n E and h E (1..... H2) being the index of the head.
........
A
A
A
A
X
A
X
(22.2)
X
2
X
X
(2.33
X
X
+
.
4
-
-1
-1
42
(13)
The attention matrices (212.1)
(2,H4l) can be formalized as:
(1, -1, otherwise,
irijio & and (i - j) mod K = h
A (2
where the condition i j A ensures that all entries in the first & rows and columns are set to -X and imposes lower triangular structure due causal masking. The modulo operation instead assigns each diagonal multiple of K to +} (allowing attention) and the remaining diagonals to -À (masking attention), while h determines the shift of the the first positive diagonal to ensure the heads do not overlap. The output of each head in the second layer is given by computing:
[[e,, k(1 ((2.0) [le,, egl, hg9
((2.M)
Xsign(A (22,)))
27

---

three attention heads (arrows different colors) copies non-overlapping Published as a conference paper at ICLR 2025
Psc
AL1
ALL
14.7
14.3
PER
.09.3 .0231
Figure 19: Visualization of the mechanism of the second attention layer in our construction. The ma- trix represents the input of the second layer (2) whereas the single vectors the output for the 10th token 1,12.1) A19.2) (/(23). Each transition probabilities at distance 3 from each other from the past. By doing this, the output of the second layer for the current token (10) contains all for each lag without loss of information.
and applying softmax and in the limit as À 0o, the rows of the attention become uniform for positive entries and zero otherwise:
A (Softmax(A
1 i h+) mod K ==
[m
mod K = 0]
which in turns implies that the attention computes (.(2) = Aun(h A(2.)),
= (())) Ej= A (2.1)
h
The output of each head is then concatenated into the residual stream. The visualization in Fig. 19 is helpful in fully understanding the mechanism implemented in the second layer. Due to space constraints, we only show the output for the 10th token for each attention head A((2.), (((22) and
k,223) We can see how each attention head, represented by the arrows of different colors, aggregates the transition probabilities by attending to the tokens the past that do not overlap and averaging them with uniform weights. It is important to remember that the input of the second attention layer is given by the concatenation of the input and output of the first layer, (1) = [ex:: cs R. thus when the tokens are averaged by the attention, the output will also contains an average of the position e;
and semantics es, of the attended tokens which are represented by the first half of the vectors (2.2A)) Interestingly the block summing the positional encoding is basically storing a copy of the attention row A!2)) in the token i. Finally, when concatenating the output of the difference heads, can see how the 10th-token will store the transition probabilities for each lag and of its entire past.
Third layer, average of transition probabilities and lag selection. Once the heads in the second layer aggregate the normalized transition probabilities from previous steps and store them within each token's 's embedding, the third layer sums these probabilities and uses the result to determine the correct lag that generated the sequence.
To build some intuition, let's revisit our running example. Suppose the current token is at position i = 10, and are predicting the 11th token in the sequence. Given set of possible lags {1, 2, 3), the third attention mechanism must focus on one the tokens at positions 8, 9, or 10. This ensures that transitions for all possible lags are considered: If the sequence was generated from the source of lag 1, the token at position 10 needs to be copied to predict the transition probabilities for the 1 th token. For lag 2, the token at position 9 is copied, and on. To determine the correct lag based on the sequence's history, our construction relies on the sum of past transitions up to the current token, Ejer P. Therefore, to select the correct lag, the third attention constructed so that the entries corresponding the transitions of possible lags are proportional to the respective
28

---

Published as a conference paper at ICLR 2025
cumulative sums. For example, to select which token among the ones in position 8, 9, 10 should be copied to predict 11, the third attention must be such that A is proportional exclusively to the sum of transitions of lag 1, i.e. Ejs10 PA.1 while A10,9 is exclusively proportional to Σjg P3,2 and A X ii P3.3. Then, the limit the softmax converging to the hardmax, the attention collapses the entry corresponding the larger sum and select the correspondent lag by copying the associated token. More generally, for this to apply to all rows, the third attention matrix must be
constructed such that A! x Ejs: Ps.1, while (3)
(3) Ejsi and A to Ejsi P73. with
all remaining entries zero.
This selection mechanism is implemented through the combination of multiple blocks within the third attention matrix, A which acts on the concatenated tokens n!²)
[n] (1) 1 (2.1)
(.(2 and is structured as follows:
400
(33
A
B 3
with the general formulation of the two matrices (3) and ((3) given by:
- [1 if ffijttee ii-jjiikK*
+1, if (i-j) mod K= = K -
B = 0, otherwise.
The matrix (3)) acts on the positional embedding of the input, similarly the matrix A in the first layer. The difference is that the position of the diagonals is now shifted by one. This shift ensures that the only non-zero entries after softmax are the ones on the diagonals corresponding to k-1 for K. The matrix (3) is instead responsible for the sum of the normalized transitions. Each block operates on the output of corresponding head in the second layer. To understand how,
consider the following tokens output of the first head in the second layer,
P4.a
94.1
Pr.a
fr.2
fr..
Prous
Pie.a
P1o.1
1/5
0
O
0
0
Pb.5
Pou2
Pa.
Pb.s
Fs.2
Ps.1
(2)b) abilitics and m. (2.1)
where we define pi
(2,b) which contains the normalized transition prob- E R2 the block of hig E R the block that contains the copy of the second attention. By the structure
of (3) and the concatenation of the outputs (2) [h(0), ((()) (2221) ,..., (2(2222)) we can see how
(33) act on these two blocks:
K Σ m h=1
(2) (2)
p (2.)) B(3) (2,h) + e
where each operation involving ((3) is selectively summing the transition probabilities from the correspondent head. As previously mentioned, for the sum to be selective, it must hold that EjsiPaa where we replace with i - k + because, after applying soft- max, these will be the only non-zero entries due to (3) The key idea behind this operation that ((3.k)m(2A)) m is boolean vector, such that when multiplied by .(2.)) it sums only the entries that correspond to transitions of lag k. For instance, consider the product Pio ((,1)) ((3
1! (2)T (2))
(2,1) (2.1) in Eq. (14), it should sum the transitions of lag 3 stored in (2.1) to give A1o,s such that the 8th token can be copied to predict the 11 th if the lag of the sequence is 3. However,
notice how simply taking
the inner product ((2.1) mg
T (2.4) would lead to the wrong computations summing over the transitions
29

---

Published as a conference paper at ICLR 2025
of lag 2 instead of 3 and excluding the transition correspondent to P1,2- This happens because all the transitions are stored starting from the element
of pe p(2.1) and not i. To account for this, the matrix B(3,1) performs permutation such that the mask is shifted by one position. Along with permuting the mask, the matrix ((3.1) also removes the normalization factor (1/2) and includes the missing transitions in the mask. To achieve this, each column of (3.h) follows pattern in which the entries are spaced at intervals of K, and the pattern shifts by position between successive columns such that all possible sequences are present allowing to sum over all possible lags. This shift creates eyclic arrangement where the columns repeat every K the transitions within the vector ((2.k)
Raaa
0
Pa.1
p 1221)T B(3.1)
mg
(2.1)
Pra
0
0
3
Phi
3
0
0
3
(14)
Fio.a
P1o,3
Prou
0
F10.1
Pro.1
8
P10.3
P7.3
+
P4.3)
More in general, the matrix (3) takes the following form:
B () 3 +1, if ((i- -icjj-K+ie K+ mod K = 0)
(15)
otherwise
The additional blocks containing ((3) act on outputs the other heads, performing the same operation by selectively summing the transitions of the same lag stored in the respective outputs. Considering heads and only the non-zero entries after softmax, occurring at j = + due to A33), we get
K
K
The
(1) (
k+1+
Σ
P3
(2
B
(2,)) m 1+11
Th
Pithtnk, + +
Am1
n=D
where Th = - Applying the softmax taking the limit X oo results in non-zero entries of the attention only where A +λ similarly to the first layer. Moreover, for large for which Th +1 i-f and absorbing K inside the temperature 3, considering the last token T we recover the weights in Eq. (3):
exp
(T h)
Pn,n-s41
amk+1
if T T-j+1=k for kE (16)
A (
r
exp
7T
fay
elsewhere
G.1 SELECTIVE INDUCTION HEAD AND NEXT TOKEN PREDICTION
Selective induction head. Eq. (16) shows how the last attention layer computes weighted average of the tokens at distance & from the next one, with weights proportional to the average of the normalized transition probabilities lag k. In practice, trained models often learn large values of 3. Thus, we consider the limit 3 oo, where the softmax converges to the hardmax:
33 = 1 [i -1 k*] with h= ====aa argmax
Σ Pa.k-
Here, the transformer selects the causal structure (i.e., the lag) corresponding to the largest En PjAk- Given the current token i, the third layer then copies the token from the position
30

---

Published as a conference paper at ICLR 2025
i k- 1, i.e., n,00) -i
11
4++1 After concatenation to the
residual stream, the tokens of the following form:
h Cs (11) (221)
(0) -- 1-k
(2.10
+1 -
1-
Output layer: next token prediction. Finally, the output layer Wo RSx2, contains all zero blocks, except for the one acting on the semantics of the token copied by the third attention. This block learns the transition matrix P to predict the transition probabilities the next token via
Wo = (0sxs OsxT Osxdo 0sx2do
Osx2d p.T OsxT
Osxdo Osx2do
Osx2do)
i... Woh (3) p* 68._ P s This final layer shows how transformers can learn selective induction head. This mechanism adapts the input sequence by copying the token corresponding to the argmax of specific quantity extracted by the previous layers and stored in the embeddings.
H ALTERNATIVE THIRD LAYER CONSTRUCTION USING POSITIONAL EMBEDDING
In this section, we illustrate an alternative but equivalent construction that implements the same predictor in Proposition 1. The first and second layers remain identical, the only difference is in the third layer which implements the selective sum of the normalized transition probabilities. This selection mechanism implemented through the combination of multiple blocks within the third attention matrix, (3)), which, in this alternative construction is structured as follows:
0
(99
A
0
(17)
0
B(3.))
We notice how, compared to the construction in Section 4, the blocks B(..1, ,-.., (3,
now positioned all in the first column. Moreover, they are not parameterized by the same matrix contrary to the other construction. The matrix (3) acts on the positional embedding the input similarly to the matrix (2) in the first layer as the previous construction:
4
1
-
3
+1 if j-i+lek
-1
+1
-1
-1
-2
-
A
A
[-1 j-i+lgK
12
1A
1A
+1
+2
44
4
4A
4
4
4
This ensures that the only non-zero entries after softmax will the ones on the diagonals corre- sponding the lags seen during training. The matrices B(3.1)... B(,, are again responsible for the summation; each matrix operates on the output of corresponding head in the second layer.
To understand how this selective sum is implemented, let us consider the output of the first head in
the second layer h (2)
[[e.
el.
A
(1)
(2.1)
h (2.Ha) in our example for the tokens 8.9 and 10:
k/2.4)
Cau
Puo
P1.1
Pr.a
Fr.
P7.1
Pio.a
Prc.
Pro.1
Cas Co
Pu.a
66.9
Pr..
Ps.a
31

---

Published as a conference paper at ICLR 2025
and define p3 (2.A) R the block of ii which contains the normalized transition probabilities.
By the structure in Eq. (17) we can see how, when computing the attention, the matrices (3,A) act on these two blocks:
(2) h 1
(2)
E h=1
Pi ( B(a + c
here we notice how the difference compared to the construction in Section 4 lies in the fact that, due the position we are not using the copy of the attention to construct the boolean vector but directly the one hot encoding of the position. Each operation involving (33,k) is still selectively summing the transition probabilities from the correspondent head but with a slightly different mechanism. Let us consider the product ( [1 which will be the only non zero entries after softmax, and show how it only sums the transitions of lag k. The main idea is that (3. are boolean matrices such that each column sums only the entries containing the transitions for one of the lags. To achieve this, each column in the matrix follows pattern in which the entries are spaced at intervals of K, and the pattern shifts by one position between successive columns. This shift creates a cyclic arrangement across the columns which repeat with frequency K. For each head h, the matrix (3, is structured such that the product 2(22k) (33 results a vector where each element the sum of the transitions for given In particular, the first element of the vector corresponds to the of kmin- the K -tth element corresponds to the sum of kmax. and this pattern repeats cyclically for subsequent elements. .(2.1) (33,)) in Eq. (I.2), which sums the transitions To give an example, consider the product Pie
stored in (2.11), Notice the structure of (33.1); the first column aligns with the transitions of lag in Pio Given that the index of p. (2,1) is 10 and the index of e, is 8, the sum constructs A(10,s) which ((.11) is used to copy the 8th token to predict the I 11 if the lag of the sequence is 3. Hence, P10 B (3,1))
has to sum the transitions of lag 3:
P4.3
PP4.1
Pr.1
P1o.1
P4.1
P4.2
Pr.2
P1o.2
(2.1)T P1o
B(3.1)
eg
8
Pr.2
D
0
0
P4.5
Pr.a
P10.a
0
3
3
Pioa
Pr.1
Pr.
Pho.1
0
0
Pro.a
P1a
Pri
Pho.a
Pra
P1o.1/
9
P10.3
P7.3
+
P4.3
3
We can see how the operation implemented by this different parameterization is the same then in the other construction. Therefore the overall predictor remains unchanged. The additional matrices (33,),, which act on the outputs of the other heads (2.1), perform the same operation by summing the transitions stored in the outputs of the respective heads. The difference in the construction of the matrix ((3, for h + 1 is that the columns are shifted by h positions relative to h = 1. Specifically, for each h. the columns are shifted by / positions compared to the matrix (3,11) In more generality. the matrix ((3,)) is constructed as follows:
B(a 3 {+1 if (i-j-h +1) 1) mod K 0)
(0.
otherwise
where h takes into account for the shift.
I CONSTRUCTION FOR ANY SET OF LAGS
The construction illustrated in Section 4 considers only contiguous lags, i.e. set of lags that are intervals of the positive integers. However, both our interleaved Markov chain framework and the Transformer construction can be extended any set of lags. The implemented algorithm is the same, but the structure of the weights in the different layers becomes more complex because the mechanism with which the transition probabilities aggregated depends the relative distance
32

---

Published as a conference paper at ICLR 2025
between the lags in the set. Due to the difficulties in finding a general formulation of the matrices involved for any set of lags as well as the optimal number of heads which depends now not only on the number of lags but the relative distance between them, we limit this section into illustrate two example with T = 10 and K = {1,3} and = 12 K = (1,3.4} for which we will visualize the matrices and operations involved.
1.1 EXAMPLE FOR K. = (1,3)
First layer: The structure of the first layer remains unchanged from Section 4. The important difference is that now the diagonals in the matrix A()) with positive entries are only 2nd and 4th
A
A
log P
0
0
A
A
-
10
+A if j-icK
4
4
4
if --i
4
4
4
The output token at index i after the first layer still corresponds to a weighted average of the past tokens h for k E K. where the weights are given by the normalized probabilities pt,
8 Attm Attn(1ig. (1 A()); =
E -i-
jEK
Erek
Payar Paniss
h2
if > 1
11
if
ph ifi k
Σ
i-k
(0)
ifi=
Due to the lack of the entries on the 3rd diagonal, both the attention and the output token will change accordingly:
0
A4
PLR
PSX
PL1
PS.1
0
0
0
0
11
0
0
0
(1)
B
0
1/5
0
0
0
0
01
0
0
11
0
0
0
0
0
0
0
(18)
Second layer. Similarly to the construction for contiguous lags, the second layer responsible for aggregating the normalized transition probabilities such that they are stored in the embedding of the current vector for its entire history. The second attention needs to learn an effective way of doing convex combination of the input tokens such that the overlap is minimized and all the transitions are stored without mixing them. Consider the token i in Eq. (26), summing two consecutive tokens such as fi and ( ,contrary to the contiguous case in Eq. (12), does not lead to any mixing due to the absence transitions of lag 2. Therefore, 2 attention heads are still sufficient to copy all the transitions in the past as long they learn to attend two consecutive tokens each.
Therefore, the optimal way to combine past tokens strictly depends on the number of tokens and the relative distance between them. Hence, finding general formula for the positions at which the second attention (2)) should be attended to minimize overlap, is challenging and beyond the scope of this work. Similar considerations apply to the optimal number of heads required, which depends on the solution of the previous problem. However, the task for arbitrary sets of lags, always be solved by consider the correspondent contiguous problem with & min(KC + heads. However there are cases in which we can leverage the structure given by the distance between the lags to use fewer heads. One example is the one considered in this section with A. (1,3) only need two
33

---

Published as a conference paper at ICLR 2025
heads to achieve optimal sample complexity. The form of the matrix ((2 remains unchanged:
A
0
(2.1
Considering the case illustrated in Eq. (26), in order for the two heads to copy all the tokens without overlap, is sufficient to sum two consecutive tokens and skip two. Therefore, the first attention has the pattem: (0, 0, 1, while the second one (1, 1, 0, 0) as illustrated in the following:
A
A
A
-
A
A
A
-
-1
-
-
-1
-
-
-
-
-
-
1
-
-A
+
-A
-
(2.22)
-
-
-
-
-1
-1
-4
-
-A
(19)
-
A
1
-
4
41
3
4
4
3
44
47
where the first i rows and columns empty because the first i elements of the sequence sampled independently from the stationary distribution and therefore no transitions are present.
The attention computes the operation before:
1 (2.h)
-
h Atn (1 A(ah)) ((2)))
Σ
1
(2,1)
+1
The output of each head is then concatenated into the residual stream. The structure of the third layer for the general case of any set of lags, also needs some modifications to take into account the particular structure that enforced in the second layer. We extend the construction introduced in Section H using the positional embeddings. First of all. the matrix A remains unchanged compared to the previous constructions, it has positive values along the diagonals correspondent to the lags shifted by one position to take into account the fact that are predicting the next token in the sequence:
1 = +X j-itlex if j-i+lEx
A
(20)
4
The matrix (3) is responsible for the sum of the normalized transitions; each block operates on the output of corresponding head in the second layer. To understand how, consider the following tokens in output of the first head in the second layer:
1/4
Σ
Ps.s
Ph.,5
Pb.1
Po.1
Pu.s
Pun.
P1.1
Pau,1
1/4
C.
&
Pa.i
Ps.s
P4,1
Pb.
Ps.
Ph.ss
9h,
fh,1
1/4
Σ Car
Pa.a
P4,1
Pr.a
Ps.
Fr.1
F%,1
(21)
By the structure of (3) applicd on the positional encoding ej and the result multiplied by p!
how, when computing the attention, the matrices (3,A) are
((2.))
K
(2) 1 (2)
Σ
(2.h) p
B(3.)), + (3)
(22)
34

---

Published as a conference paper at ICLR 2025
where ((3,A) is selectively summing the transition probabilities from the correspondent head. As for the simpler case of contiguous lags, for the sum to be selective must hold that h! \(3) f- x Ejz; Pa.ce where i k + are the only non-zero entries due to A after applying softmax. As before, B(3,)) are boolean matrices such that each column sums only the entries containing the transitions for one of the lags. To achieve this, the matrix need to learn the same patter in the attention of the second layer 2A which was used to sum the vectors and create the current inputs. Each column is shifted by one position and they eyclically repeat with frequency K. In the following example, we consider ((2.1) g(3.1e eg in Eq. (23), which sums the transitions stored in (2.11) in the entry Aloss
0
0
0
U
Pa.a
Pa.a
Paa
Piaa
3
Ps.1
B
Pa.1
(2.1) es
A611
0
D
4
P9.
Pa.
(23)
Pan.8
0
PLn.a
0
Po.1
P10.1
Pro.1
0
3
P10.3
P9.3
P6.3
P5.3
Notice how the matrix (3,1)) has the same pattern as (2.1) in Eq. (19) but along the columns instead of the rows. Intuitively it makes sense since need to the same entries resulting from the sum in the previous attention. The matrix (3,2) acting second head will have the same pattern but shifted by two positions in order to have the same pattern as (2.2).
1.2 EXAMPLE WITH (1,3,4)
The case of two lags K. {1, 3) despite not being contiguous does adequately represent the general case. Indeed due to the structure, we could always sum two consecutive tokens and therefore recover optimal performance using heads. It is helpful to also consider case where the lags do not form structure that allows for fewer heads in the construction. For example the case of three lags = (1,3, 4} and 12:
First layer: The main structure of the first layer remains unchanged, the diagonals in the matrix AO) with positive entries are 2nd 3rd and thh
A0
log PT
0 (1)
-1
A0
4
4
4
(24)
+λ if j-iEK
4A
-X if j j-i4K.
47
The output token at index i after the first layer still corresponds to weighted average of the past tokens ( for is e where the weights are given by the normalized probabilities PL.::
0
10
14
3
3
0
P PM
(4
(Riz)
59
3
3
0
0
0
0
SP
(43)
PA
0
(RALA
(25)
35

---

Published as a conference paper at ICLR 2025
Second layer, aggregation of transition probabilities: In this case can't use anymore the fact that 2 consecutive tokens can be summed without mixing the information. In fact, summing the last tow tokens h and hig would now result in the mixing of P11,3 and P12,4. In order to avoid this, the only possibility to sum one token every similar to the case where we would have contiguous lags. This solution is less efficient because summing each 4 tokens while having the missing transition corresponding to lag leaves an empty element in the embedding of the token and adds additional head, increasing both the dimension and the number of parameters. This means that even we only have lags, in order to not have any overlap we still need 4 attention heads for construction not mix the information. Each head has the pattern (0,0,0, 0, 1) shifted by one position as if the lags would be 1,2,3,4:
A
-X
4
-
A
2 2
-
A(2.))
A
(
4
4
A22,3)
X
(2.44)
A
4
4
4
4
4
Third layer For the third layer we use again the construction with the positional encoding that was introduced in App H:
1
A +1 if j-i+lex
A
-
if j-i+1gK
For the selective sum, the matrices (3.1)) ..... (3,4 have the same form before but considering now the fact that even we only have 3 lags in the set, we still need heads:
+1, if ((i (i-j - mod & min(X) + 1 1=00
B (3.4)
1o. otherwise
The computation related to the matrix (3311) in the attention are reported in the following:
0
0
O
T
g(3.1)
P12
10
2
0
D
2
P1n.a
Pio.1
0
Piz.1
P1x.1
P12.3
P8.3
2
36

---

Published as a conference paper at ICLR 2025
J CONSTRUCTION FOR TWO LAGS AND SINGLE HEAD
In the constructions illustrated so far, in order to store all the transitions the history of the current token and not lose any information, we had to scale the number of heads at least as the number of lags the task K. This allows to achieve optimal sample complexity. However, driven by experimental evidence, we observed that scaling the number of heads the number of lags is not necessary in the special case of |K] 2. In this case indeed, there exists a solution, which transformers can learn, that achieves optimal sample complexity using only one head in the second layer. In the following we will report the construction that proves the previous statement while illustrating for the case of X }1 analogous to Section
First layer: The structure of the first layer remains unchanged from Section 4. The important difference is that now the diagonals in the matrix A()) with positive entries are only 2nd and 4th
A
log P
0
0
A
2
40
10 +λ if j-ieK
4
4
-1
- if j-i4K.
4
Remarking that each input element is encoded as 4, = [es.: es] E {0, 1] SH+ the output token
index after the first layer still corresponds a weighted average of the past tokens 11 for k E IC where the weights are given by the normalized probabilities Pt.::
i(") (1
Σj 1
li j [EK] K]
Frek Prr. h2
Prp
(0)
if 1
1
ifi=1
Σ
Picch
ift> 1
kER
h
ifi=1
Due to the lack of the entries on the 3rd diagonal, both the attention and the output token will change accordingly:
0
16
.4(1)
PL
PL1
P5.1
0
0
0
1
1
0
0
0
0
(1)1)
0
0
0
0
0
0
0
Bea
0
0
0
0
(26)
Second Layer:
the second layer uses only the first head A(2) = ( compared to the construc-
tion illustrated in Section and the matrix (2) = 1221) remains identical.
4
A
-
-1
4
-
A 2
4
-
A
-
42
A
A
4 4
0
2
4
-
X
+1
For the case of two lags examined here, we can derive mathematical expression for the matrix A(2) which is valid for any set of lags. For convenience we introduce is min
0, if < k or < k
1
X if < and (ii (パi- mod 2(k -k)) < (k- R)),
(27)
0, otherwise
37

---

Published as a conference paper at ICLR 2025
where the first condition ensures that all elements in the first l rows and the first columns of the matrix are zero. The condition j I i ensures that only the lower triangular part of the matrix (including the diagonal). Finally, the condition (li - jl mod 2d) < d introduces periodic pattern within the lower triangular part of the matrix. The modulo operation creates a repeating cycle of length 2(k - k), and the condition < (k k) determines whether to place a one or a zero within each cycle segment.
Third layer: the third layer instead has a different structure. As before, there are only two non- zero blocks A and B(3) but the latter appears in the transpose position compared the previous constructions:
(
0
(
A
B
(28)
The matrix (3) remains unchanged and has positive entries along the diagonals, corresponding to the lags shifted by one position. The main difference lies in the matrix B(3), which now includes negative entries in positions that previously contained zeros. The general formulation this matrix is the following:
0, if
1/
+ if jiitlek
B
+3, if ((i jii-j-1) mod 2 k)) (k: k),
if j-itlgk
8, otherwise.
So far the matrix (3) has been structured such that it would compute the selective sum of the nor-
malized transition the lag of the corresponding entry the attention: A(3) x (2)) / k-+1 x Ejss Pj.ke where i & are the only non-zero entries due to (3) after applying softmax. To understand the impact of having negative entries, let us consider the previous example for the case of KC (1,3) and the output of the second attention the 8th 9th and 10th token:
1/4
0
0
0
#5.3
16.3
Ps.1
$6.1
Pin.5
P9.1
P10.1
1/4
6
0
0
0
S: 8
P4.a
Fs.8
P4,1
As.
As
PR.1
Pb,1
R
1/a
P4,1
Pr.a
0
and define (2) E RT as the block of (2(2) which contains the normalized transition probabilities
such that h EjjN, Cxys m()) Ejen, 1 By the different structure in Eq. (28) we
how, when computing the attention for the concatenated tokens h,2)) [[e.,> ex], i(()) A((2), the order of the multiplication has been reversed (e; is now the left) and the matrix ( is applied to
(2).
p
1OT ( e, B(0) + (
(29)
To better understand the implications of the reverse order in the multiplication and the presence of negative entries, consider the product To (33.11);;(2) in Eq. (31), which sums the transitions stored in
38

---

Published as a conference paper at ICLR 2025
((222)) which, after softmax, will correspond to A
(3)
10,8=
0
UN
PL
0
0
0
0
T B()) (2)T 3
0
B
0
PrL1
(30)
0
P8.3
P7.3
P4.3
P8.1
P7.1
P4.1
where we observe how, the product involving ( now not only computing the sum of transitions for the lag 3 as for the previous constructions to copy the 8th to predict the 11 it is also subtracting all the transitions of lag To fully understand the implications, we also consider the entry of the attention correspondent to the other lag in the set, 1 and the relative product ejo ((311) (2).
0
0
0
0
Psa
Ps.a
0
0
0
0
PL.1
Ps.
100
T B(4) (2)T
0
0
0
0
Pau
Pb.3
0
0
Pan.3
0
0
P1o.1
P1o.1
3
110.1
P9,1
P6,1
P5.1
P10.3
Pg
P6.3
P5,3
(31)
therefore both products contains the sum of the transitions for the respective lags and the negative sum of the other lag and notice how they are all computed on different elements of the past. The first one contains the transitions for the tokens 8,7,4 whereas the second one contains remaining ones 10,9,6,5. If we compute the softmax:
exp
6
B(3)
T
P10
+1)
A10,10
Σ 1= exp
c
B(3
+ exp
< B(3) (2)
+2) + exp
cT B(3 +
(TOB()))10))
exp
exp
B(3) p/
21 22) + exp
B() (222 + exp
((T0(3)pp9))
Considering the limit of \ 20:
B() Pio
exp
lim A10.10
A
e B() (2)
c10 B()
exp
+ exp
1
T ((3) // - B(3))pp)
-
+ EiE Pca - Eie Pi, E (i Pr.1 + Σ (10,000,0555 pr,3 +
exp
which considering all the possible transitions as for the case two heads therefore achieving optimal sample complexity.
39


Reviews:
[{"review_id": "9CSANwPCYX", "reviewer": "Reviewer_acym", "summary": "This paper proposes a new synthetic task based on interleaved Markov Chains in order to better mimic the complex token interactions in natural language. The authors come up with an algorithm to solve this task with a 3 layer attention-only (disentangled) transfomer. Notably, the authors also show that when they train a similar model with gradient descent, it learns similar mechanisms (apparent by comparing the attention maps).", "questions": "* The synthetic task of modeling interleaved Markov Chains is interesting and definitely a step in the right direction. But do you think it is yet a good enough proxy to model the complex dynamics of token interactions? In language modeling, often these interactions are informed by the semantics of the tokens and the context. Whereas I don't think here the token meanings are important in any way.\n    \n* In my opinion, one of the things that made the induction head work interesting is that they found heads behaving in similar ways in actual large language models which were trained on real textual data. I am not really sure if the same can be said about the selective induction head circuit you found here. If this is something that can only be observed in a synthetic task, do we really care about it? Or, do you have any hypotheses about how selective induction heads might manifest in models trained on natural language texts? And, how do you plan to investigate this hypothesis (maybe in future work)?\n\n* In Section 5 you mentioned that a model when trained with less than k attention heads per layer was also able to successfully solve this task. This either suggests a completely different algorithm or is indicative of head-level superposition of information. Can you get more insights into this by looking at the attention maps?\n    * Also, have you tried with more than K heads per layer? I wonder if it would further break the mechanism. Or, you might start seeing duplicate heads behaving in similar ways.\n    * A similar question about the number of layers --- You need at least 2 layers to form the usual induction head behavior. Do you really need at least 3 layers for \"selective\" induction heads? Also, what happens with more layers? Do more layers introduce more \"steps\"? Or, some of the layers essentially become identity functions?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* The paper is well written, I learned something.\n* Introduces a new synthetic task and discusses how transformers can successfully solve it with an interpretable mechanism. The authors also show that the same mechanism is learned when training with gradient descent, which I found most interesting.", "weaknesses": "* I have no major concerns apart from the utility of this insight in the broader language (or sequence) modeling tasks that we actually care about. See questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LhVVGRFfD4", "reviewer": "Reviewer_o5dv", "summary": "This paper presents an in context learning task that requires a model to infer what type of markov chain was used to generate the input tokens in order to predict the next token. Then it constructs a solution to this task with a three layer transformer that involves a particular type of attention head they call selective induction heads. In the final two pages an experiment is reported on where a three layer model is trained on this task and then the attention head patterns are visualized. The attention pattern in the constructed model and the trained model is highly similar, suggesting that the very same solution is achieved in both cases. \n\nThe paper is well written, and great care is taken in the presentation of the technical details to the theoretical solution. However, the task here is pretty idiosyncratic and a specialized transformer architecture is used to make the analysis easier to perform. The connection to actual in-context learning scenarios seems like it might be a bit fraught, I’m not sure how much closer a lagging markov chain is to natural language than a normal markov chain. I also get the feeling that moving to a standard transformer architecture would make things a lot less pretty...\n\nGenerally, I am amenable to technical papers that deal with toy settings, so I think the core content of this paper is solid. However, I don’t think that the title and framing of the paper is accurate. The process of selecting causal structures in context is incredibly general and complex, and selective induction heads are just a step in the direction of understanding how in context learning occurs. I think the title and framing needs more transparently communicate that this work is largely theoretical. How you do this is totally up to you and I’m willing to raise my score if you address this concern.\n\nThe attention pattern visualization is remarkably similar between the constructed model and the trained model, so while this analysis is correlational, I would expect the causal claims to follow as well. That being said, you might wonder whether a standard transformer model trained on this task implements basically the same solution (I assume that the attention head pattern is different, otherwise you would have included results). I think the right way to understand this involves causality, and all the work that you reference comes from an AI safety mechanistic interpretability background, but there has been loads of work done in causal analysis of neural networks. Here are some review/history papers to check out and get some citations from:\n\nhttps://arxiv.org/abs/2408.01416\n\nhttps://arxiv.org/abs/2301.04709\n\nhttps://arxiv.org/abs/2410.09087", "questions": "See review", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "See review", "weaknesses": "See review", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WzSShKoErg", "reviewer": "Reviewer_7pgv", "summary": "This work introduces a new circuit in Transformers, “selective induction heads” that enables transformers to determine the lag in causally-structured data as per the context (“in-context”) and apply the induction head copying mechanism. A 3-layer disentangled Transformer model is constructed to achieve this, with specifically designed purposes for each layer towards identifying the correct lag: capturing the transition probabilities, aggregating the probabilities, and summing the probabilities to select the lag. Empirical validation on synthetic data shows that the constructed transformer can achieve performance close to the theoretical maximum likelihood estimator across settings (contiguous, non-contiguous, many lags, single head), and the attention maps verify that the layers are performing the functions they were designed for.", "questions": "* Continuing from the note on scalability in the “weaknesses” section: is it necessarily true that there is such an idea as a fixed “correct lag” that ever appears in real data? This is partially addressed in the many lags and non-contiguous lags setting, but appears to require further generalization for such a transformer model to be truly useful (for example, in natural language tasks).\n* Are there any core (architectural) considerations that might need to be made for scaling to larger data, from an efficiency standpoint? How many heads would be needed, and can this be reduced? This is addressed to some extent in the many lags setting in an Appendices G and H.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The general setting is defined in a clear manner, and intuitively demonstrated through Figure 1. \n* The 3-layer disentangled transformer perspective is an interesting construction that is well motivated by designing specific purposes for each layer; the preliminaries are also explained quite well. \n* The paper is written clearly despite it being admittedly challenging to follow some of the notation — the visualizations are immensely helpful to follow the construction in Section 4, as well as the various special cases examined in Appendices E, G, and H. These appendix sections were greatly appreciated, as they thoroughly flesh out the nuances in construction for each such setting.\n* The attention maps serve to reinforce the model construction; each layer serves their respective intended purposes.", "weaknesses": "* More discussion in the empirical validation would be useful, to explain the findings and walk through the plots in Figure 4. For instance, it would be beneficial to explain the difference between the 1 head (1H) and 3 heads (3H) cases, as the proposed construction scales the heads with the number of lags, and detail the implications of these differences. \n* The scalability of the method remains unexplored; that is, what are the empirically-realized sets of lags that appear in real-world tasks, and how do they vary across tasks? Many natural language reasoning tasks have a non-linear and more complex causal structure, and it would be good to see how / whether this method extends to such tasks, beyond synthetic ones. \n* While Proposition 1 is defended by construction (and the empirical findings), it is unclear how Claim 1 is supported empirically (although it is acknowledged that this claim has yet to be formally proven) through the results in Section 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iHj6DkJmkR", "reviewer": "Reviewer_DFwY", "summary": "The paper mainly extends the work of Nichani et al. (2024). Nichani et al. show that Transformers can learn causal structures of Markov chains and estimate transitional probabilities via In-Context Learning (ICL). The current works investigates whether Transformers can learn causal structures from a more complicated variant of that task - particularly from interleaved Markov chain, where a prediction at position t, is not dependent on the state a t-1 necessarily, but some state at t-k. So in this context, Transformers have to not only learn the transition probabilities and Markov chain dependencies, but also have to learn to select the right Markov chain for prediction at t ignoring other interleaved ones that are not relevant for prediction at t. \n\nThe paper makes a theoretical construction based on Disentangled Transformers to show how Transformers, in principle, can utilize in-context examples to solve the task. The paper then empirically investigates and show that Transformers can learn variants of the task through gradient descent.", "questions": "n/a", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A reasonable extension of a prior publication (Nichani et al. (2024)). \n1. Provides deeper insight into ICL capabilities of Transformers. \n1. The task could be another sanity check for alternative architectures to test whether they can do what Transformers can or not.", "weaknesses": "1. The scope of impact may be a bit limited. While the experiments and construction do deepen our insights to ICL, they don't seem to provide a fundamental shift in our understanding, and may not translate into being particularly informative for specific practical decisions. But this is a more subjective point and I am not factoring this too strongly for the evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iGahAglFd9", "reviewer": "Reviewer_Lmxm", "summary": "This paper studies the ability of transformers to identify the causal dependency in Interleaved Markov Chains. The authors start by introducing Interleaved Markov Chains as a Markov Chain with a random lag. Then, the authors give an explicit construction to show a transformer can implement a predictor for the in-context selection of interleaved Markov chains. Finally, the authors train a transformer on interleaved Markov chains and show that the transformer not only solves the detection task but also exhibits similar attention maps to the theoretical construction given by the authors.", "questions": "* Do you have any recommendations for an LLM to improve its understanding and learning of causal structures?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper provides strong evidence that transformers can identify causal structures in the context.\n* It is nice to see that the theoretical construction can match the experimental findings in a controlled experiment.\n* Some slightly more general cases are discussed: Single-head transformers & Non-contiguous lags.", "weaknesses": "* The scope of this study is a bit narrow. The authors only discussed the detection of interleaved Markov chains\n* It is unclear whether an LLM can detect interleaved Markov chains. Maybe the authors can do some inferences on open-source models to verify this.\n* It remains uncertain whether an LLM can retain its ability to detect interleaved Markov chains when trained on data that is unrelated or only loosely related to such structures. Given the artificial nature of this task, it's more effective to assess detection ability in an out-of-distribution setting rather than through direct training on the same task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)