How Transformers Learn Causal Structures In-Context: Explainable Mechanism Meets Theoretical Guarantee
Authors: Jianzhe Wei
Siyu Chen
Jianliang He
Zhuoran Yang
Venue: ICLR 2026 Poster (2026)
Keywords: transformers
in-context learning
interpretability
Markov chain
Primary area: foundation or frontier models, including LLMs
TL;DR: We show that transformers can and do approximate Bayesian model averaging to infer varying causal dependencies in-context, with information-theoretic guarantees.

Abstract: Transformers have demonstrated remarkable in-context learning abilities, adapting to new tasks from just a few examples without parameter updates. However, theoretical understanding of this phenomenon typically assumes fixed dependency structures, while real-world sequences exhibit flexible, context-dependent relationships. We address this gap by investigating whether transformers can learn causal structures -- the underlying dependencies between sequence elements -- directly from in-context examples. We propose a novel framework using Markov chains with randomly sampled causal dependencies, where transformers must infer which tokens depend on which predecessors to make accurate predictions. Our key contributions are threefold: (1) We prove that a two-layer transformer with relative positional embeddings can implement Bayesian Model Averaging (BMA), the optimal statistical algorithm for causal structure inference; (2) Through extensive experiments and parameter-level analysis, we demonstrate that transformers trained on this task approximate BMA, with attention patterns directly reflecting the inferred causal structures; (3) We provide information-theoretic guarantees showing how transformers recover causal dependencies and extend our analysis to continuous dynamical systems, revealing fundamental differences in representational requirements. Our findings bridge the gap between empirical observations of in-context learning and theoretical understanding, showing that transformers can perform sophisticated statistical inference over structural uncertainty.

Published as a conference paper at ICLR 2026
HOW TRANSFORMERS LEARN CAUSAL STRUCTURES IN-CONTEXT: EXPLAINABLE MECHANISM MEETS
THEORETICAL GUARANTEE
Jianzhe Weiz.2
Siyu Chen
Jianliang He²
Zhuoran Yang***
Georgia Institute of Technology
Yale University
jwei 355.344echhahh..... (siya.chen.ss sc322..222., fianliang he, zhuoran, zhuoran.yang syale, edu
ABSTRACT
Transformers have demonstrated remarkable in-context learning abilities, adapt- ing to new tasks from just few examples without parameter updates. However, theoretical understanding of this phenomenon typically assumes fixed dependency structures, while real-world sequences exhibit flexible, context-dependent rela- tionships. We address this gap by investigating whether transformers can learn causal structures - the underlying dependencies between sequence elements - di- rectly from in-context examples. We propose novel framework using Markov chains with randomly sampled causal dependencies, where transformers must in- fer which tokens depend on which predecessors to make accurate predictions. Our key contributions are threefold: (1) We prove that a two-layer transformer with rel- ative positional embeddings can implement Bayesian Model Averaging (BMA), the optimal statistical algorithm for causal structure inference; (2) Through exten- sive experiments and parameter-level analysis, we demonstrate that transformers trained this task approximate BMA, with attention patterns directly reflecting the inferred causal structures; (3) We provide information-thewn -heoorettccccc guarantees showing how transformers recover causal dependencies and extend analysis continuous dynamical systems, revealing fundamental differences in representa- tional requirements. Our findings bridge the gap between empirical observations of in-context learning and theoretical understanding, showing that transformers can perform sophisticated statistical inference over structural uncertainty.
INTRODUCTION
Modern transformers exhibit a remarkable capability: they can adapt to entirely new tasks using only handful of examples, without any parameter updates. This phenomenon, known in-context learning (ICL.) Brown et al. (2020), has revolutionized our understanding of what neural networks achieve. A model trained on diverse text suddenly perform arithmetic, translate languages, or write code all by simply observing a few demonstrations. Yet despite extensive empirical success Wei et al. (2022); Garg et al. (2022) and theoretical investigations Von Oswald et al. (2023); Akyürek et al. (2023); Goel & Bartlett (2024), fundamental question remains: how do transformers adapt to the varying dependency structures present in real-worth sequences? (Allen-Zhu & Li, 2025; Bietti et al., 2023; Zhao et al., 2023; Wibisono & Wang, 2024)
The Theory-Practie Gap. Current theoretical understanding ICL rests on a critical simplifica- tion: most analyses assume that dependencies between sequence elements follow fixed, predeter- mined structure, For instance, theoretical works typically study settings where tokens are indepen- dent [[<1, fz [z2, (x2)), Chen et al., 2024a; Wang et al., 2025). However, natural language and real-world sequences exhibit far richer structure words depend on previous words in complex, context-dependent ways that vary across sentences and domains. Recent work by Nichani et al. (2024) began addressing this by showing transformers encode fixed causal structures during training. Specifically, they assume
follow rigid patterns like (x1), J2, f(2)] (Bai et al., 2023;
"Corresponding author.

---

Published as a conference paper at ICLR 2026
an n-gram causal model (e.g., bigrams where each token depends only the previous one) (Ra- jaraman et al., 2024: Edelman et al., 2024), and prove that transformers can embed this structure in their attention weights to perform inference. However, in real-world scenarios, the dependency graph itself is not fixed but varies across different sequences. For example, in language, the syntactic structure can change dramatically between different documents, and in stock price prediction, the relationships between assets can shift over time. Thus, a key challenge is
Can transformers infer and adapt to causal structure in-context?
Our Approach. We introduce a novel framework where sequences are generated from Markov chains with randomly sampled causal dependencies. In our setting, each token depends exactly one predecessor, or its "parent", but crucially, these parent relationships are not fixed and must be inferred from context examples, which is a collection of sequences sharing the same underlying causal structure. This setup captures the essence of (x) by requiring the model to adapt to different latent structures across contexts. The transformer must infer these latent dependencies from con- text examples to accurately predict new sequences - mirroring how language models must adapt to different syntactic structures or reasoning patterns.
Main Contributions. We consider two types of Markov chains: discrete chains over a finite vocab- ulary and continuous linear dynamical systems. Our work makes the following contributions:
For discrete Markov chains, we prove that a two-layer transformer with relative position embeddings can implement Bayesian Model Averaging (BMA), the statisti- cally optimal algorithm for inferring causal structures from observations. Our construction shows how attention mechanisms can perform sophisticated probabilistic inference over structural uncer- tainty. (2) Empirical Verification: Through extensive experiments on Markov chains, we demon- strate that transformers trained via gradient descent converge to solutions remarkably similar to our theoretical construction. Parameter-level analysis reveals that learned attention patterns directly encode posterior probabilities over causal structures, providing mechanistic insight into how trans- formers perform statistical inference. (3) Information-
(1) Theoretical Construction:
Analysis: We establish conditions under which causal structures can be recovered in-context, using mutual information and data pro- cessing inequalities. Additionally, we show that gradient-based learning naturally discovers these structures early in training through x² -muuull information maximization. (4) Extensions to Contin- uous Systems: We extend our framework to linear dynamical systems in continuous space, reveal- ing fundamental differences in how transformers handle discrete versus continuous causal inference. While transformers show strong empirical performance, we identify representational limitations that prevent exact BMA implementation in continuous settings.
2
PRELIMINARY
2.1 TASK SETUP
To investigate the question (*), we consider data generated from distributions with latent causal structure. Each sample is a sequence of tokens T1.H = [æ1...., æH.. where the h-th token æy depends on one of its predecessors, called the parent token ppa(h)- This dependency relation is represented directed tree graph
(NE()) where pa(h) ~ Unif(1,.... h 1), vh {2.... H}. Given the causal structure defined above, the data generation process be writ- ten as Ty = G(Tpa())). where G(-) denotes either stochastic sampling from the transition kernel T( (+()pa()) of Markov chains, or the autoregressive process for dynamical systems. G(-) is fixed during the sampling of the whole dataset.
For the in-context learning task, suppose we have L+1 samples (w( [e[[+11] ( from the same causal graph G. The first L samples are provided as in in-context demonstrations from which the model infers the latent graph structure, while the last sample is the target for prediction. Except for the first token 2/++1, every token x I++1 in this trajectory is required to be predicted via next-token prediction conditioned on x 1:L. 1:H and its past observations L++ 1:h-1'
Markov Chain. Following the Markovian assumption adopted in Edelman et al. (2024); Nichani et al. (2024); Chen et al. (2024b), we assume that sequences are sampled from a Markov chain with random dependencies. In this setting, tokens {æa) are drawn from finite vocabulary V
2

---

Published as a conference paper at ICLR 2026
Example
Example
Example L++
Task overview of in-context causal structure learning. Each training sequence consists of I. examples with observed variables and hidden parent relations, followed by new example 1++1 where the model must infer the underlying parent indices in context from previous demonstrations.
Figure 1:
fei,..., ed), where [V] = d and {ei} are one-hot vectors. The random dependencies are specified by latent causal graph (nE()))) Let V A(V) denote the Markov transition kernel, where A(V) is the probability simplex over V. Then each token is generated as 24, x(- pa(h)) A(V), wh E [H], where, by slight abuse of notation, we also regard и the stochastic matrix E R MIXVV with rft 5l m())) E, Tī. 1.
Dynamical System. Beyond the discrete Markov chain case, we also consider more challenging setting with a continuous sampling space. Here tokens {ær) are dense vectors in Id. The link function g()) replaces the discrete transition kernel, and we instantiate it as linear dynamical system with additive Gaussian noise: 2h g (Fa(())) pA pa((h) / p² mh; where A E ((R)) is orthogonal, 1 (0, La). Th (0, Ia) and ensures the stability of the sequence.
These settings evaluate the extent to which transformers can perform in-context causality learning.
Goal: Inferring the Causal Structure. The task formulation naturally raises the following ques- tion: Given in-context examples, how can the model infer the underlying graph structure G? A classical approach this problem is Bayesian Model Averaging (BMA), which leverages Bayes' rule to compute the posterior distribution over the possible parameter space. Treating the parent structure pa(h) as the parameter be estimated, the distribution of having parent ht will be pre- dicted as its posterior probability given L observations:
P( pa(h) = h/) P(pa(h) h²) En"e(m P P( pa(h) = h") P(pa(h) = h²)
P(pa(h) = a
(I)
By Eq. (1) and our task assumption, we have the following lemma of the formulation of BMA.
Lemma 1. Suppose L samples are observed from the Markov chain (or dynamical system) X1:If with latent causal structure g. Bayesian Model Averaging makes prediction of pa(h) E 1):
exp(EIe[1) log 7 (x)1aa)))))
Phpa(h) = 1/' |x(:ft)
(p ( (log
(2)
Enee(h-1] exp exp Eiel, log ক (2t
where gr (log x) E RR-11,
h.L PM
EiEt2 log π (ar (x//)) See Appendix C.2 for detailed proof.
This Bayesian formulation provides principled baseline for inferring causal structure, and serves as point of comparison for the in-context learning behavior of transformers.
2.2
MODEL ARCHITECTURE
2.2.
STANDARD TRANSFORMER
Decoder-only Transformer is neural network structure for handling sequential data. Given quence of tokens (wi...., WT), transformers first embed these tokens and add positional encod- ing: A E(w.) P(L) Rd, VI [T]. In matrix form, the mapped tokens of the input represented as ()) = 1 Txd Subsequent layers consist
headee attention layers (MHA) followed by multilayer perceptron layers (MLP). At layer I, the hidden features (1-1) are updated as follows. First, the causal-mas mask self-attenie attention layer computes the ouput by:
Attn(H; Wo, WK. Wy) = =o(M
HWQ))
HWv,
ddz
3

---

Published as a conference paper at ICLR 2026
2nd Attention Layer of Parent Selection
1et Multi-lfead Actention Layer
Head
Pawm Cancliess
Head
Head
invetcenc tune
Cant
Pane Cancblone
Figure 2: Visualization of Attention Weights A A(2), Left: A 6 examples. Transformer shows the capability to select true causal tokens (red rectangle). Right: Six representative heads (out of 10) from A (Head and 10 degenerate). Trained with 10, H 10, 5, 1024 steps.
where We, WK. Wv E Rdxds o(V); E, expley) applied to matrix row-wisely, and M is the causal mask where M(X); -00 ifi> else Xij. Then multi-headed attention gives the output:
exp(v:)
M + m=1
MHA(H) =
Attn(H W W W Wo:
where + denotes the concatenation of vectors and Wo E RMdeedd After obtaining the interme- diate features M from the attention layer, this feature will be added to the residual stream, which aggregates the previous output: H' (1-1) MHA The FFN layer adopts this input and updates this stream as:
FFN o ((WW) W2, H A (H)
where W1 E Rdxdm. W2 E IRd xd and (-)) is the activation function. Finally, the output of the L-layer Transformer is a H (w)),, projected to vocabulary logits by Wo RdxV
2.2.2 DISENTANGLED TRANSFORMERS
To better analyze the role that each part of transformers plays in learning task, prior works Fried- man et al. (2023) propose the disentangled transformer, which decouples the intertwined features in the residual stream. Instead of adding each layer's output, the disentangled transformer concente- nates it with the residual stream. Suppose transformers, the hidden states are ( d'dd-
HO (H (--1) Attn (I")), Attn (H [TX(++M)4)
Here we consider the decoder-based attention-only transformers. And following Nichani et al. (2024), in each attention head, Wx W reparameterized by WxQ. Wo by Wov and the initial input (0) is given by h [E(w,,, P(w,)] [ew, ed] E Rd+T,, In our task, the input sequence consists of L + 1 examples of length-H H chains, leading to the positional embedding size 10, then d 10 < 110 in the input embedding
set d H and WKQ in the first layer will have 1 @(10*) parameters. Instead, considering w, as the h-th token in example 1, we use two types of embeddings to represent this positional informa- e, E R PosH(w,) en " This reduces the required parameters to O(H² L²) for training. The formulation of this transformer is based on Eq. (60) and (69).
equal (L + 1) H)
tion: Posı
Relative Positional Embedding. While the original transformer employs the absolute positional embeddings, subsequent research has demonstrated the advantages relative positional embed- dings (RPE) (Shaw et al., 2018; Su et al., 2024). To enable tractable parameter-level analysis, adopt a simplified structure based on RPE that reduces the parameter space. Crucially, empirical re- sults presented in Appendices G, H and demonstrate this simplification didn't prevent mechanistic
4

---

Published as a conference paper at ICLR 2026
Cross-Fxample
Exaen-e RPE w!
1eead of Cww.OCCCCOe RTC
Relative Poation
thay pard chidiiin
InWWeol of 48m
(Transition
Now Indenne
238
Figure 3: Parameter visualization of 2-layer transformer with RPE. Heads of inner-example RPE
w uniformly show the largest value at position h 0 except Head 1, 9, 10. Correspondingly, W KQ shows similar blocks diagonal except block 1, 9, 10. Besides, o(Wov) approximates
interpretation on more standard architectures like standard transformers with FFNs and disentangled ones. Instead, the following simplified model helps to understand the mechanism of transformers. For RPE, is commonly parameterized by vector E TR. which assigns attention score w only via relative distance
between positions of query i and key j. Recall the sequence length = HLL + 1). Similar to the absolute positional ones, we adopt two types of RPE: Wt E g representing the order from L+1 examples and w E 222H denoting the order h from H tokens.
(h( w M], (h, E [H]?, W2(,t') = (wLl 11), 1> I'
-0, else. (for causal mask)
In the transformers we consider, the first layer adopts this RPE. Its output u, for token æ, = ar, is:
Attilan -R1T
Σ oe (w (h, + wr((,,))aa
CE[T
(4)
(xp(wH(( (h, h²) + (L(1)))
Σ
(' Ee exp(wn WH h²) + wL(1, We
Suppose have K heads in the first layer. The outputs {u. kEjK] will be concatenated by disen- tangled residual as the input zt for the next layer. Then the next single-headed self-attention layer takes features last example I++1 as query, key and value tokens and gives the final prediction. Recall the input is क1.7 = x 11L+1 This transformer architecture is formulated as follows:
1st RPE Attention (K-head):
u = Attn
-11. = <w (h. + w (L 12
E Rd
Disentangled Residual:
[1]). ... K
2h = | L+1 v..]
E Rd
2nd Attention (1-head):
fir( (H())
(zI Sith-WWKeH W WKQZz) T IT
=I Wov T
=o(uTA- W Wkquh)
xish1 Wov
Rd
(5)
where fte(- (7(b) E Rd denotes the output of the transformer based on context H/h [al: 11hh-1 (or the context denoted by 7/ for brevity), and we assume some blocks in WKQ. Wor are 0:
Oaxd
Oax Wkq
Woy 0axкa 0kaxa
WkQ
Wov
Oxaxd
0gaхка
where W KO E Ka Wor Iddd are trainable and represented by WKQ and Woy for brevity in the rest of the paper. To train transformers, cross entropy loss is used for the Markov chain (MC) and MSE loss for the dynamical system (DS) shown in Appendix Eq. (13).
5

---

For attention layer, the attention weights A normalized by the reveal which tokens query primarily attends to, enabling mechanistic interpretability analyses such as circuit discovery Olsson et al. (2022). We first investigate the attention patterns AA A2(2) from the first and second trans- former layers. Mathematically, they matrices where the i-th row denotes the attention weights Published as a conference paper at ICLR 2026
1st RPE Layer
Masked
2nd Artention Layer
Masked
xT
AD
Bled
Hleod 2
Head
x1
xg_
x1
x11
Copier: xb
Kon
Queer
23 42*
XA
(w PMPP (b.) w/ 1.)
x.) 10
War
Figure 4: Overview the constructed mechanism in transformers: the st layer works as copier, and the attention patterns 2nd layer follows BMA, which approximately select correct parent token.
3 CAN TRANSFORMERS IN-ONNEXT LEARN CAUSAL STRUCTURES?
3.1 2-LAYER TRANSFORMER LEARNED TO SELECT CAUSAL STRUCTURE IN-CONTEXT
To investigate the question (*), first train 2-1ayer transformers with RPE introduced above on the Markov chain setting. Each input has L+ 1 samples {æ of Length-// Markov chain with causal structure G. The input sequence length T H (L- + 1). We set the transformer have K heads in the first RPE layer: {(w w [E[K] and head for the 2nd attention layer (WKQ. Wov). All RPE parameters are initialized randomly from Gaussian distribution and (WKQ: Wov) from zero.
of token æz to the whole sequence and Al = A is formulated by the following for i, E [T] and ( E Rxx A(2) E HH
A (MM( 1
h2) + ().. A = V/ WKQ
where the index attributed to æ; from Eq. (5) which is the h2-th token of 1,-th example (simi- larly for index j), and Ux(') are the hidden features VIH of the L + 1-th example from Layer 1. Empirically, we observe the trained attention patterns of (A(2) match the groundtruth causal struc- ture in Fig. 2. For the first layer, shown Fig. some heads of attention weights show specific attendance among tokens while some heads degenerate (e.g., Head 1, 9, 10). Further, we visualize the trainable parameters of 2-layer transformer w w WKQ, Wov. Positional pattern in w diagonal pattern in WKQ and the similarity between Wov and log ग can be observed in Fig. 3. To fully understand why the transformer can select causal structure, we analyze theoretically.
Takeaway 1. Transformer formulated by Eq. (5) effectively identifies latent causal parents in- context (Fig. 2) and learns highly structural parameters aligned with the task (Fig. 3).
3.2 CONSTRUCTED TRANSFORMERS IMPLEMENT STATISTICAL ALGORITHM
Based the patterns observed in experiments (Fig. 3), make the following assumptions for the transformer defined by Eq. (5):
+1,
h=0,
3 E [4] s.t. whll = 3 +1, 1= K, 1-1, (
w [4]
-1, he [=] \0,
W
Oaxa
Oaxa
Oaxd
W
Oaxd
WkQ
a (Wov) = n
Oaxd
Oaxd
w

---

Published as a conference paper at ICLR 2026
where W is unknown parameter and for RPE, assume that 0-th entry dominates wfr and one element k' of w dominates it: w H [0] > w [-0, w wilk) [K]
> wi[ - Since the K heads are identical up to their indices, we assume without loss of generality that the dominant entry of w occurs at position k. i.e., w [ 3 and we set K = L. The theorem below shows that, aligned with the above restriction, a constructed transformer can implement statistical algorithm for inferring the causal structure {pa(h) hidden behind æ i.1 and predicting 2/+11 T(- (x pa(k))
Theorem 1. Under the restriction by Eq. (7), the transformer fo is parameterized by O {(3, W)).
Then for fa with W log π in Eq. (7), its second attention layer A (H; approximates Bayesian Model Averaging (see Lemma 1):
lim A A (H,0) A4000
lim (.. (P
(8)
B-+00
Further, the transformer's prediction of the last example EI.H with L context examples converges to the true conditional distribution given the causal parent guaranteed by Theorem 2:
lim fol- H/) = π(- ppa(k)), wh [H]. 8.
Proof Sketch. Figure 4 gives an overview of the construction: in the first RPE attention layer, each head from Eq. (7) is assigned to retrieve one historical copy of the same token Th. so that concatenat- ing L heads recovers past L observations 3 In the second layer, with the condition in Eq. (7), the attention score between tokens (h, h') reduces to bilinear form Ph (W) E, ay W ah. which by W log ι coincides with the BMA score Рana = Eie(z log π ( la/e)) With the causal mask, the softmax attention exactly matches the parent-selection distribution in BMA. By the theoretical guarantee of causal token selection (Theorem 2), OV matrix Wov receives the correct parent an L (a))) and makes prediction for 7(- 2 [ pa(h)) The full technical proof is deferred to Appendix C.I.
D' Angelo et al. (2025) also considers an in-context causal learning task. With minor modification of the above construction, we can show that transformers can implement BMA for that easier task.'
Takeaway 2. Two-layer transformers can explicitly implement B.M for causal token selection.
3.3 WHAT ALGORITHM DOES THE TRANSFORMER LEARN?
Although we have constructed a transformer implementing the BMA algorithm, what do transform- ers actually learn after training? Since the core lies in the attention weight (4 with WXQ which recovers graph structures, we next analyze its characteristics in detail. In the following, we use Wer to denote the trainable submatrix in Eq. (7). We first define the parent selection metric from cross-cntropy loss, which quantitively shows the accuracy of models to predict parent indices:
(A (a1H (x 1-++1 ().9)
1
Σ KE[H]
(TO(A) log A
1
Σ KE(H)
log log. (2) (2)
(10)
H.
H
where (2) is considered as an algorithm for predicting parent (pa(h) given the input #1. 11++1
and we have (ABN,,)) = H Ene(m (ps (log by Eq. (2) where (W) =
Σ, =T Wz/. We visualize this metric Epo during the transformer training process in Fig. and compare it with BMA's. We observe that the transformer's parent selection loss decreases dur- ing training while remaining above the loss of BMA, gradually approaching it.
Generalized parent selection with varying size L'. We further test how well the transformer and BMA generalize in parent selection under different sample sizes L': since A(2) and Asna are formulated via p² Eie(e) iish-1 War,, we vary the number of demonstrations as set of L', (W) with W (Wil) logm).
and finally compute p² A/.' and the parent selection loss LL' CL'(W) pa
A comprehensive comparison with D' Angelo et al. (2025), highlighting both similarities and distinctions, included in Appendix A (Related Work).
7

---

Published as a conference paper at ICLR 2026
Larter Bair
Figure 5: Generalization of parent selection loss ( for transformers trained with L E = 10, and H = 15 with first layer fixed as constructed.
(1. (1)....20)..... d
20),
From Fig. 5, observe that: 1) across different test sizes L', the trained transformers achieve performance close to BMA (loss differences mostly within a small margin); 2) models with smaller training length L generalizes better, with loss curves approaching BMA more closely; 3) for model with fixed training size L, the parent loss decreases rapidly as L' increases, converging toward zero. The above results show that the trained transformers have comparable performance BMA.
Parameter Verification. Beyond behavioral agreement, a crucial question is: can we interpret the parameters of transformers with the aligned BMA inference? We evaluate the similarity between the trained weight W and the theoretical BMA parameter W = log π. As row sum of π equals 1. we first check whether o (Wer)
म. Note that, oW.r)) = π <> W.r log + 61 where 61 denotes a row-wise shift canceled out by the row-softmax a. This softmax of (W::) provides reasonable way to normalize KQ matrix making its scale comparable
ग with scale [0, 1]. However, the empirical results in Fig. 6 (first three subfigures) show that misalignments exist between or (W=+) and иτ. From the view of the output prediction, we can see the attention mechanism ( Wo.) introduces a degree of freedom on the column of log rather than the row:
Proposition 1 (Invariance). the columns of W differ from those of log т by column-wite shift, i.e, WT = log ग + 1a .Ya E Rd, then we have the same output of attention module and BMA:
o
Wesch) = =o(E, 1:--1111111 log nr²),)
Further, if Markov chain æA.H is stationary and W log π + la + T1 a. b E Rd, the above conclusion also holds asymptotically as - DC. See the detailed proof in Appendix C.3.
This proposition indicates that instead of checking whether o(W.r) = a(log (). we should check whether e (.+) cal (log iT) as verifies if Wt = log क + 1a T. We follow this evaluate the discrepancy between a (W=z) and aea (log r). As illustrated in Fig. 6, the deviation remains small, with column softmax error
- |o (Wrr) clo T))l < 0.05.
(11)
This alignment also holds across various training model sizes d E (10, 30, 50), as shown in Fig. 10. confirms that trained transformers did approximate log ग + la wwhich is equivalent to implement BMA method via Proposition 1. Taken together with theoretical construction and empirical results, we conclude that transformers can and do implement BMA for in-context causal structure learning.
Takeaway 3. Transformers with trainable W.r closely approximate BMA in causal token selec- tion (Fig. 5) and learn parameters which explicitly implement BMA method (Fig. 6).
Robustness to Model Architecture. A natural question might arise: is the specific RPE-based construction essential for this mechanism? We show in Appendix G that standard disentangled transformers with absolute positional embeddings are also theoretically capable of implementing BMA. Empirically, we trained such disentangled models and standard transformers with FFNs (de- tails in Appendices G-D). Despite the increased parameter complexity, these models converge to the same attention patterns and achieve parent selection performance comparable to BMA (Fig. 19). This confirms that the interpreted parent selection mechanism a general solution found by the optimization process, independent of the simplified RPE parameterization.
8

---

Published as a conference paper at ICLR 2026
COMIATON
Figure 6: Parameter-level comparison between transformer W and BMA log т. Here, Wtr de- notes the diagonal block being trained. Trained with d 20, H 50, L 3, and 2048 training steps while Wov remained fixed. (See results with Wov and Wer both trainable in Fig. 9.)
3.4 THEORETICAL GUARANTEE OF LEARNED ALGORITHM
Beyond identifying what algorithm a trainable transformer adopts, we further establish the theoreti- cal understanding of why transformers can select the correct causal token via information-theores principles. Our approach follows Nichani et al. (2024), which leverages Mutual Information to- gether with the data's 's inherent property of the Data Processing Inequality (DPI), In contrast to their gradient-based proof, we show that transformers can exploit this property directly in context. More- over, our analysis generalizes the 2--uul information framework of Nichani et al. (2024) to the setting reducible to classical mutual information, by exploiting the cetemattionthhoreeti structure characterized in Lemma 3 and 4. Finally, our proof of Theorem 2 applies to finite-horizon Markov chains whose marginal distributions not assumed be stationary.
We leverage the data processing inequality regarding classical mutual information 7 and -m-uull information Lx² (see Definition 1 in Appendix C.10) to establish the identifiability of causal struc- tures. I, xx can be uniformly derived from f-divergence which helps to prove DPI for generalized f-mutual information y. These information metrics reveal an essential property in data:
a
Lemma 2 (DPI. Theorem 3.9 and 7.16 in Polyanskiy & Wu (2023)). If random variables X + + 2, i.e., satisfy the Markov property p(x, y, 2) = p(z))(y)(x) p(( ly), then we have Iy(y; 2 1 2). Further, for classical mutual information, I (x; z) = 1 (y; z) iff I (x: = 0 iff + y.
In Markov chain setting with random causal structures, X Xpa(h) X holds, which guarantees the DPL. Further, following Nichani et al. (2024), we develop a stricter version of DPI for classical mutual information in our setting, guaranteeing the selection of unique parents.
Lemma Let X1:H be a Markov chain with transition kernel w and causal structure G. Sup- pose there exist ", E (0, 1) such that min (Ts'ls) 7/IV| and the marginal distribution min, P(xx- = s) > Then there exists a S 1- by < such that: (p(t)))
((x);x2)
AN' < h, h + pa(h).
The aforementioned assumption is the transition lower bound condition in Nichani et al. (2024), which ensures uniform minorization of the transition kernel and implies the strong DPI; see Ap- pendix C.4 for details. Building on the above lemma, we establish the following result.
Lemma 4. For a Markov chain X1.H with causal structure G satisfying the conditions of Lemma 3, suppose further that I [xxi Xpa(k)) > 0. Then
Ex [log T (Xx [pa(k))] > Ex (xxl(x)))))] Vh' * pa(h).
This establishes the DPI in the expected log-liielioood (MLE) form. See Appendix C.5 for details.
Then the following theorem shows that the attention weights of transformers leverage in-context ex- amples to instantiate the above criterion for causal structure. The proof is deferred to Appendix C.6. Theorem 2. Under the conditions of Lemma 4, consider the transformer constructed in Theorem 1, which implements the BMA method. The attention weights Al..
=A( (mt:h) satisfy:
lim Af. lim (M( 1 1
Cpa(h) H where Phe
Σ log T (2%
Takeaway 4. Information-theoreds analysis reveals that the parent selection exploits the condi- tional entropy, where the strong DPI guarantees the identifiability of the true causal parent.
In other words, the guarantee does not require samples to be drawn from stationary distribution and depends only on the finite-horizon transition structure.

---

Published as a conference paper at ICLR 2026
3.5 CAUSAL STRUCTURE IN TRAINING DYNAMICS
We further investigate the training dynamics of the transformer. We show that random causal struc- tures embedded in the inputs will be recovered the gradients of loss w.r.t. the core WKQ matrix:
Theorem 3 (Informal). Consider the transformer So constructed as in Theorem with trainable diagonal block W of WKQ specified in Eq. (7) and trained with cross-entropy loss
((0) Ei Ex
Ex,g [og (a L- -17) c)] +c]=== Eg ((00 9)
(12)
where E(0: h. 9) x Dog fo(
Let P denote the intermediate parameter p(W)
En 2 W a from attention scores. If the Markov chain is stationary. i.e,, Br =²= vh [H] and the initial feo with W O outputs μ for any input, then the gradient at initialization statisfies:
de h,
oe(Oo; h,
wh' + ppa(h).
ap
0p
po(h)
In derivation of the gradients, we show that these terms are highly related to x2 information, which reveals the numerical relation in gradients via data processing inequality. See Appendix C.7 for the detailed proof. This result explains how transformers extract meaningful structural infor- mation from data via intermediate parameters, rather than directly encoding this information to its parameters. Empirically, Fig. 11 shows the gradient of Bl(0)) matches the latent causal structure.
Op
Takeaway 5. Gradient at initialization is able recover the latent causal structure, driven by x²- mutual information, which facilitates structural discovery in early training (Thm. and Fig. 11).
4 DYNAMICAL SYSTEM EXTENSION: FROM DISCRETE TO CONTINUOUS
Further, we investigate the Markov chain in continuous space, where we examine the linear dynami- cal system with latent causal structures: X pAA (pa(h) + /1 p² m. E Rdd, 17, (00 Ta). We initially train a transformer with RPE introduced in Eq. (5) on data generated from the dynamical system. Similar experiment results on attention weights (11), AA))) and parameter visualizations can be found in Appendix Fig. 12, 13, 14, and 15. Visualization of RPE parameters is consistent with the construction in Eq. (7). Moreover, the attention weights (2) of the transformer yield accurate predictions of parent indices across various examples. Similar to the discrete case, we can define the transition p(-l-) by Xn (po(h) ~ (pA Xpa(h)) (1 ) Consequently, Eq. (2) similarly specifies the BMA formulation under the dynamical system setting. In this context, Lemma 4 remains valid and guarantees the asymptotic correctness of BMA's 's parent selection. To investigate transformers' mechanism of parent selection, we test the parent selection loss C/2s of the transformer and BMA in the dynamical setting, where we set varying L' in-context samples introduced in Sec. 3.3. Fig. 17 demonstrates that the transformer with trainable (W, Wov) achieves performance comparable to the BMA method when L' approaches 20, while the loss CLds remains with a noticeable gap as L' is small. We conjecture that the proposition below explains this discrepancy:
(Representation Limitation of Transformers). Under the assumption Eq. (7), both the transformer and BMA take the unified form An-th = (p"")w w. In the DS setting, transformer logits are bilinear, p'r, = E, ar Wer whereas BMA logits are Pl = E, T. Aat, + d E larhe 1|2 with d + O. Then there exists no W such that a (p't) = o (Pben) holds for all DS samples (=1++11 ,G) ~ Ps. Hence, transformers under Eq. (7) cannot represent BMA in the DS setting. However, in the MC setting, Wer = log n matches the BMA formulation.
Proposition
See Appendix C.8 for detailed proof. For BMA, (Dbm)) ) = E, log p(af, ( In dynamical systems, transition p(-l-) involves not only cross but also quadratic terms. When the representation equation holds, substituting 24, by PA Ppa(h) + VI /I-a0 p²² mp. and using the independence of mh. coefficients of T) must vanish, which requires W = ci A. But this Wer fails to represent BMA's quadratic term. So no W et can yield A a (Mip" (W:))) as BMA's.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transtormers learn implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neu- ral Information Processing Systems, 2023. URL https: L2 LziniAXEI9, 14
/ /tpp///veeeeeeeeeeeeeeeeeeeeeeee net forum?id
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in in-context learning? investigations with linear models. In The Eleventh International
Conference on Learning Representations, 2023, URLhttps: http///pppee//eeeeeeeeeeeeeeeeee net / forum?
id==000x488N441.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part learning hierarchical lan- guage structures. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL
https:://pperevve/ net///ewweeeeeeeeeeeeeeeeeeeeeee www.mrii-nnnnni 2oruuriddmPQKKyZAAK 1,
Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mei. Transformers statisticians:
Provable in-context learning with in-context algorithm selection. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttttss///openeeeeee /tpp///pppeeeeeeeeeeee/eeeeeeeee net forum?
id-liMSqUuVge
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: memory viewpoint. In Thinth-seventh Conference on Neural Information Process- ing Systems, 2023. URL httppps:///opeeveeee.eeeeee hwtp::/ee//ooommm. d- 2-tt//oooom
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Hemighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learn- ers. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020. URL http:://prroceedngg...t neurips cc sacp/pppr__ftee Paparrappr.... pdf.
Siyu Chen, Heejune Sheen, Tianhao Wang. and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learnnnng Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442. 2024a.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. In The Thirty-eighth Annual Confer- ence Neural Information Processing Systems, 2024b. URL /tsttp:://ppeereeeee forum? 4= NN2REs OMa. 2, 14
Francesco Angelo, Francesco Croce, and Nicolas Flammarion. Selective induction heads: How transformers select causal structures in context. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https:/ http///pnnneeeeeeeeeeeeeee net. www.eeetwwffreeeeeoeeeoee.i bnJgzAQjwf. 7,
https:///ppeneeriww.ntttt
Ezra Edelman, Nikolaos Tsilivis, Benjamin L. Edelman, Eran Malach, and Surbhi Goel. The evo- lution statistical induction heads: In-context learning markov chains. In Advances in Neural Information Processing Systems, volume 37, pp. 64273-64311, 2024. 14
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https://
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. In Proceedings of the 36th International Conference Neural Information Processing Systems, 2022. ISBN 978171387 1088.
11

---

Published as a conference paper at ICLR 2026
Gautam Goel and Peter Bartlett. Can transformer represent Kalman filter? In Proceedings of the 6th Annual Learning for Dynamics & Control Conference, volume 242 of Proceedings of Machine Learning Research, pp. 1502-1512. PMLR, 15-17 Jul 2024. URL https: httppp/::///ccceeengg.... molr.press. v24 goe 1 2 html.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with repre- sentations. In The Twelfth International Conference on Learning Representations, 2024. URL net/ derrum?idiiwwwevaaa...
14
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Inter- national Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL 24
worg/abs///..oooo /112/ 6980.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A curious case of single-layer transformers. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//openreview nwt. foott//foummat/ f=t 4 14
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pp. 18018-38070, 2024. 1, 2, 4, 9, 14, 15, 19, 21, 37
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. 1895, 2022. 6,
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2024. URL https: / openeeviewww....////e net. Seeee//formm..ddkS55dddaas..
14
Yury Polyanskiy and Yihong Wu, Information Theory: From Coding to Learning. Cambridge University Press, 2023. ISBN 9781 178130822,0 doi:
9
Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vardhan Makkuva. Transformers on markov data: Constant depth suffices. In Advances in Neural Infor- mation Processing Systems, volume 37, pp. 137521-137556, 2024. doi: 10.52202/079011-43699
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Pupers), pp- 464-468, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10118653/1/N11/ 2074. URL ptatpp://aclanthologg..o wrt.loogoorrggg/l/e wwt///8-///00 2011--2074..2
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. RoFormer: En- hanced transformer with rotary position embedding. Neurocopuut, 568(C), February 2024. ISSN 0925-2312. doi: 10.101 www.::.ccccmm.....2222.. neucom. 2023 127063.
URL https: :tapps://ddd/ /doii org /c.. 1016/j.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174. PMLR, 23-29 Jul 2023. URL https:/ https://proceedi http//pocceenngge/// .ll...r...... malr.press/ v2222 wow///ooooooaalllloooooooomll www.wowwwllll.oo................ d23 html. 1, 14, 29
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers can learn tem- poral difference methods for in-context reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025. URL https: http///ppee/eeeeeeeeeeee net / forum? id= 3o 6mxCxP1.
12

---

Published as a conference paper at ICLR 2026
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V Le. Finetuned language models are zero-shot learners. In Interna- tional Conference Learning Representations, 2022. URL https::/ net/ forum? i AGCozdgr.
Kevin Christian Wibisono and Yixin Wang. From unstructured data to in-context learning: Ex- ploring what tasks can be learned and when. The Thirty-eightt Annual Conference Neu- ral Information Processing Systems, 2024. URL https: / net net/forummidt x9eF 9aVVRI. 1, 15
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference Learning Representations, 2022. URL https:
RAUVEEE 14, 15
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(4)):1-55, 2024. 29
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization. In Proceed- ings of The 28th International Conference on Artificial Intelligence and Statistics, volume 258 of Proceedings of Machine Learning Research, pp. 1684-1692. PMLR, 03-05 May 2025. URL https: httt::::/pxocceddnggs
malr.press/ / v228/2252///// zhang25d.htm. 14
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? In Proceedings of the 2023 Conference on Empirical Meth- ods in Natural Language Processing, pp. 16513-16542, Singapore, December 2023. Associa- tion for Computational Linguistics. doi: 10.18653/w1/2033 / /acllanthllogy... org/2023 sanlp-mainn.. 1029/.
eemnlpmmmaee matpppmaai..12.... 1029. URL https:
13

---

Published as a conference paper at ICLR 2026
A NOTATION AND RELATED WORK
Notation. We use [+] to denote the set (1,2.....,h) h). For causal structure, we use pa(h) to represent the parent index of node h. If the stationary distribution of Markov chain æ ((a)) exists, then it is denoted by μ¹ E A For the transformer, the input of a sequence of vectors is given by Ft.r == [01, 22...., TTT E RRdXT, Given the input, we denote the attention scores of a standard self-attention layer as p" == #1.T W WQ E R However the causal mask M in the attention
layer will lead to p² == x ist Wx QQ E RR.I. a(p')) = (MM(')) )),, Vt' E [t 1]. We do not distinguish between them in the proofs. For the matrix form of the attention, we use A and A to denote attention weights and scores, correspondingly, where we have p/. Ar and (M(( = A. In training, we use cross-entropy loss and MSE loss for Markov chain and dynamical system settings, respectively:
H
(M)
1
Σ
I+11
log (o(fuso Hh)) + t),
= 1
h
h=2
(13)
"
(P$
1
||
L+1 - (-to H1))1
where 0 represents all trainable parameters and e is a small value to avoid numerical issues with log.
Related Work. A growing body of work studies the in-context learning (ICL) ability of trans- formers from different perspectives. One line of work understands ICL as a form of Bayesian inference, showing how the latent concept can be approximately inferred under restrictive theo- retical assumptions (Xie et al., 2022: Zhang et al., 2025; Panwar et al., 2024). Another direction of research investigates how transformers can simulate standard algorithms, such as gradient de- scent on linear regression (Von Oswald et al., 2023; Ahn et al., 2023; Guo et al., 2024). While these works demonstrate the ICL power of transformers, they commonly assume i.idd or uncorre- lated input tokens. To move beyond i.i.d. assumptions, recent works investigate ICL with correlated data, particularly Markovian sequences (Edelman et al., 2024; Chen et al., 2024b; Makkuva et al., 2025). These settings provide insight into how transformers handle in-context learning with se- quential dependencies, but typically focus on fixed dependency structures. In contrast, our work addresses variable causal structures that differ across prompts. Pioneering this direction, Nichani et al. (2024) demonstrated that transformers can encode fixed parent-child dependencies (e.g.. bi- grams) in Markov chains. D' Angelo al. (2025) introduced selective induction heads, enabling transformers to identify the underlying Markovian order (or "lag") from a candidate set to learn this structure in-context. Our work generalizes this setting. While D' Angelo et al. (2025) focus on inferring single structural parameter (the lag k) shared across the sequence, we tackle flexible structure inference where dependencies can vary arbitrarily for each position, effectively modeling latent trees rather than fixed-lag chains. D' Angelo et al. (2025) construct a three-layer transformer that asymptotically implements maximum likelihood estimation for their task, where the construc- tion is verified via attention pattern visualization as well as KL divergence validation of next-token prediction targets. In our work, we theoretically derive two-layer architecture that explicitly imple- ments Bayesian Model Averaging (BMA) in-context. Empirically, we go beyond behavioral metrics and provide purameter-level verification, demonstrating that the trained weights directly encode the transition kernel. Furthermore, provide thcoretical understanding of in-context causal structure learning based on the Data Processing Inequality (DPI) and extend our analysis to continuous dy- namical systems, revealing representational gaps not occuring in the discrete setting.
B CONCLUSION
In this work, we investigated the capability of transformers to infer and adapt to latent causal struc- tures in-context, moving beyond the fixed dependency assumptions common in prior theoretical analyses. We proposed a novel framework based on Markov chains with randomly sampled causal dependencies, requiring the model to identify from context examples. First, we provided constructive proof that two-layer transformer with relative positional embeddings (RPE) can explicitly implement Bayesian Model Averaging (BMA). This demonstrates that the attention mechanism is theoretically capable of performing statistical
predecessor-succeess
relationships
14

---

Published as a conference paper at ICLR 2026
inference over structural uncertainty. Second, through extensive experiments and parameter-level analysis, we showed that trained transformers implement BMA method, which converge to the theoretical construction: the learned attention patterns recover the posterior probabilities of causal parents, and the weights explicitly recover the log-transition kernel of the underlying generative pro- cess. Third, we established cooommtioonthhoretii guarantees using the Data Processing Inequality (DPI), which helps understand how the selection mechanism identifies causal structures in context, and showed that gradients at initialization recover these dependencies via x²-mutual information. Finally, we extended our framework to continuous linear dynamical systems. While transformers continue to exhibit strong empirical performance in this setting, we identified the representational difference that prevents the exact implementation of BMA, unlike in the discrete case. Collectively, our findings offer a mechanistic explanation of how transformers perform in-context causal learning, highlighting their ability to act as statistical inference engines for both discrete and continuous data.
Broader Implications. Our findings support theoretical frameworks that model in-context learning as statistical inference task (Xie et al., 2022). Distinct from "Induction Heads" which typically fo- cus on copying fixed positional dependencies (Olsson et al., 2022; Nichani et al., 2024), we demon- strate a probabilistic setting where the model must infer a latent dependency structure that varies per example. This provides a mechanistic grounding for how LLMs adapt to flexible, context-dependent rules rather than relying solely on fixed n-gram statistics (Allen-Zhu & Li, 2025). Furthermore, this helps understand why LLMS demonstrate ICL capabilities on empirical tasks with "unstructured" language data (Wibisono & Wang, 2024), mirroring our setting where the transition mappings be- tween words are fixed while the structural positions of a couple of words vary from input to input.
Limitations and Future Work. We acknowledge that real-world sequences often involve complex non-linear dynamics or hierarchical dependencies (e.g., context-free grammar) beyond the Marko- vian and dynamical systems studied here. However, our primary objective in this work was to prioritize mechanistic interpretability for Markov chain or dynamical system: explicitly characteriz- ing how transformers infer latent structures in-context on these tasks. By focusing on these tractable settings, we were able to derive exact theoretical guarantees and provide parameter-level verifica- tion that the model implements Bayesian Model Averaging. We believe this explainable framework serves necessary foundation, and we leave the extension to more complex non-linear and hier- archical lata-generating processes for future exploration.
C DEFINITIONS AND PROOFS
C.I PROOF OF THEOREM 1
Proof. By the condition of (WH,WL) in Eq. (7), the attention score A of the query x, æ L+1 in the first layer is:
+1, if hhe hys le 11 ki
1 whe - he] + wi.l - [e] = 23
-1, if he + her, le - L + ki
0, otherwise,
where (e) is the (token, example) position mapping of the input æy. The output ui Attng.-aay of the first attention layer will be calculated as:
o (w (h. .) w (Z I
(14)
9 3 2 [1ph,=h.l x) )) Der L+=
where k- attention head copies the corresponding token from 4ę-th example for the query at 4+1 By using the disentangled residual, the outputs of the K heads (K" = L) will be concatenated as:
2/ [u) ...,, uk], with प. = h
I- by Eq. (14).
For the second layer, with the diagonal condition of W WxQ. the attention weight ( E HxH is given by:
L Σ
A = Un WKQU =
Wob.
= -((((((( H
(15)
1=1
15

---

Published as a conference paper at ICLR 2026
where M is the causal mask enforeing A to be strictly lower-triangular If we define the vector d E R with ph ( h- we have wh' E [h 1]:
A = -o ol p²(W) Σ - War..
(16)
where (-) is the causal mask applied to row h. setting MN(w) -00 if h' 2 V. R.. Then, we set W as log (with log applied elementwise), which leads to:
/ (log T) = A = Σ log π (2/11)//))
(17)
Considering the form in Eq. (17), Lemma 1 shows that the BMA method of Eq. (1) has the same formulation:
P(pa(h) h [m (W = log "))=
(18)
Combining Eq. (18) with the limiting behavior of the first layer in Eq. (14), we obtain the conver- gence result for parent selection as $ oc:
lim $+000
R (H: a (f¹ (W- log x)) (h(pa(h) h a 1:44)
...L We define PENa E R as the corresponding vector form for BMA whose h'-th entries are -00
if h 2 h, otherwise ph. Then we have:
B.Ih) lim (H1:0) = o (PENA)
Furthermore, guaranteed by the consistency of BMA (Theorem 2), as the sample size L the posterior estimation concentrates on the true parent pa(h) Thus, the prediction of the token distribution converges in the limit B, as:
lim p
10(
Wovv Σ I= 1+11 x(. ar )-
In proving the theorem, we rely on Lemma proved below, which illustrates the similarity between BMA and attention weights.
C.2 PROOF LEMMA
Proof. Here we use p(s(s') to denote Pa = s|aPa(h)> = s') for generality beyond discrete Markov chain. Based on Bayesian Theorem, it can be calculated by Eq. (1), Due to the Markovian property p(æ, (11-11) p(æ, pp(())), the joint distribution of this chain X1.H is:
H
p(x1.H) p(x.) II pl a 11:h-1) p(x1) II
p( ppa(k)
h==2
h=2
(19)
=p(#)) II 1
p(.: ( pp((t) pl (pppa)))
Here pa(h), pa()) in Eq. (19) are random index with prior: pa(h) Uniform (hh 1]). Conditioning on ph) = h in Eq. (1), we can substitue pa(h) (h) in Eq. (19) with h². Since {pa(1) izt are random indices out of interests, these terms eliminated:
Pa pa(h) = W))
Ih (p(æ)) Ilseen p(
ar (o())) plark, (ark,))
(20)
ΣMe Pl
l:L pa(h) = h²)
En. II. (p(a1) Hlye p(x) p(a/la/hat)) ha())) p(ar) (=))
which leads to:
cxp Elerr log p(æ)), (xhh))) Σ Eveehh-1 exp (iE(L) log pl p(z) (x/.))
P(pa(h) //' 1 (1HH)
a (1th (log W (21)
where plog.log WP) = E t T log W 2/ and the matrix W = π is induced by transition kernel P(s'(s) (T(81s) in the discrete Markov chain.
16

---

Published as a conference paper at ICLR 2026
C.3 PROOF OF PROPOSITION
Proof. First, suppose we have Wer = loo + 1a The attention scores of the transformer are:
Ptt Eath Werchh Σ x Wih-1 log ma/ + Σ x - la x/
For the second term, since { are one-hot, we have:
Σ
t -inh-la JaTT a/ -1-1a
Σ 2 hi),n
where c(a, h) = (T(22 (E, is constant with fixed index h. And by softmax operation, we have:
o (p'is) = a Σ -ish- log naze + c(a, h))
Σ a iith-1 log TT/
= oPapa (log :)),
where the first step shows that the softmax climinates the constant term, and the last equality is from Lemma 1. This shows that the transformer with We la gives the same prediction BMA's.
, vh E [H] then we can prove Further, suppose W. = log π + la T + b1 we have æn ~ p., the term from b1 also forms constant vector asymptotically:
pst
Σ
x hh- Wes
Σ
o 1hh-1 log naah Σ IT -iih-i LaT
aa
bi a
IT aTT Σ 11-11 log naH (() Σ
For each term of E, -TtT-1 ,, have:
Σ 2tfT b
L Σ Σ
Jan ss) bs
L
1 sEld
20000
b
=
Σ #E|d]
P(ow = s)b.
sE(d]
1 b,
where _.= b is a constant d(b) w.r.t. h'. Using the same technique in Theorem through division to eliminate this term which goes to infinity, we have the desired result:
lim o(pltt)- lim apte (log 7)). 1.000
C.4 PROOF OF LEMMA 3
Proof. Our target is to show there exists a < 1 s.t. 1xxx,,) S a 1 Xpa(s)) for any j < < + pa(i). Since and j are in the same tree of the causal structure G, we use p(i, to denote their lowest common ancestor (LCA).
Generally, there are two cases for the relation between node and j. as shown in Fig. 7. In the following, we first show how we can get the contraction from I(x,; X;) to [xx; Xp(t.))) ). If is not ancestor of i, then this contraction is non-trivial and from [(x; Xp(is)) to (pa(t)) we use the weak version of data processing inequality to bridge them. If is an ancestor of i, we further define the reverse transition kernel and adopt similar technique derive the contraction.
17

---

Published as a conference paper at ICLR 2026
p(i,j)
p(i,))
LCA
(Sn(s)
paff)
(a) j is not an ancestor of (branching).
(b) j is an ancestor of (ancestor chain).
Figure 7: Two cases of the relation between i and in Markov-tree. Directed arrows indicate the dependency direction (ancestor - descendant).
Recall that the definition of mutual information is
P(x: = 8,X, = ss)
[(x;xx))= = Σ P(x: s, log
(22)
p(x: s)lP(x, = 5')
Σ P(x; 5) KL(P(Xx = =P][ -), vj <i.
(23)
Let x denote the conditional probability P(Xp(i.) = -1x; = s). We can see the distribution (ii((i x is:
Mij.p(i))))) ox(5))
Σ π (()1())))))(sss)) xx(3")
(24)
Σ P(x; = s'(xp(i.)) = =s) -P(xp()))) = =s = =)
(25)
P(x, = s'ix 8),
(26)
where the last step comes from Markov property. Let μu denote the marginal distribution P(x; = .). By similar calculation, obtain:
d(..((i)))) (s))
Σ π (( (sl
(27)
Σ P(xx s' s(1xx(1.)) =s) P(xp(i))) =
(28)
P(x; = =')-
(29)
Let k ((j.p(i,.)) Hence, get the following equation for the KI. divergence term in Eq. (23):
kL(P(xx -xx; >)(le²) KL( x||m=* pl(i))).
For the term above, we apply Lemma which requires max,gv (s'(s) > 0, Vs' E satisfied by our assumption and gives the contraction for the mutual information with contractive coefficient
ag = max (|=((-)) - x-11k)]l < 11- 1- 1, as shown by Lemma This leads to:
T(xx;,)) S Σ (xxx = 8) () -II(((Xp((.)) = +))[P(xp((.))= =))
(30)
=a PP(x= s) -KL(P(xp(()) = = -Ix==)][P(>((..))) :)
(31)
aa xp(x)))
(32)
where k: k= -_jjpp(i.))) This shows the contraction from x(x;;;) to [(x;; Xp(i.j)) One may notice that the key of the proof the function of transition и on the distribution leading to the contraction in KL divergence. For bridging Xp(ts)) to I(x;; Xpe(t)), we define the reverse transition matrix to follow the above proof technique. Different from the above, since we don assume the stationary distribution of XA, the reverse transition is time-inhomogeneous.
18

---

Published as a conference paper at ICLR 2026
For the convenience, adopt to denote the distance d(i, p(i.j)) between node i and p(i,j). Then since p(i.j) is an ancestor of node i, we use (i) to denote it. Similarly, we use pa p1a to denote
the t-step parent of node i. By Bayes' rule, have:
P P(X (0) 8 (1xp'((1) FX =
((Xpa'(1) = s|xpa'- = =s)
(33)
P(Xpat t(f) 8*)
(Tssss)) [s) (Xpa = =)
(34)
P (()) s')
From the above calculation, we denote the time-/ reverse transition kernel म. Rdxd as: (s(s'') == P(Xpat(1) = s(xpat-1(0)) = =s).
(35)
Recall that the mutual information [xx; Xpa (i)) is:
Σ P(x: s') KL (P(Xpal()) -(x; = 5 8')l (Xpa((()))))) = =.))).
I [xxi Xpa' (1)
(36)
Via conditional independence of the Markov chain, we can show:
((Xpa'(0) = -(x; = 5') = .y ....... क22 P((pa(() = -|x; = 8
(37)
Suppose म.- P(Xpa(t) = =|xx; =s) P(Xpat 1(1) -|x; =6) holds. For time t:
Tt o ...o P((pa(() = =(x; = s')s = = Px((pa (() = =|xx; =k:=5 =s)ss:
(38)
= Σ P(Xpaa(i) = 8) x/xpt--1(0 = s*
= ==|xx = s'),
(39)
((Xpa'(1) (1) = 8|/;;- = =').
(40)
Similarly, we can prove:
P(Xpa' = = म, ..... .2 (Xx((()) =
(41)
From the assumption of transition kernel and marginal distribution, #. defined in Eq. (35) is valid transition kernel and we can lower bound #. by:
( (s) Ps) (xpoo(() =) 87 P P(Xpat-1(1) = s')
# (s|s')
(42)
V
With the above condition, we can apply Lemma 5 and again to get the contractive coefficient 01 < 1 δy such that:
KL(ā, o xl| v) S an KL(x]ly), VI E (4]
(43)
With these ingredients, we obtain:
1 [xxi (i)) S a 1-1 (xx; xpa(t)) a1 d(p(t.)).po()). 1 (x43 Xpa(t))
(44)
Define a max (ap, 0n }. Since 7, (), 11 by 1. And we have:
T(x;x,) S a di.n(t.5)) I(xxxx((..))
(45)
S a ((pp((t))))t((ttd..
T(x;; xpa(i))
(46)
d(ipal()) I(x;; Xpott))
(47)
5 [(xi; Xpa(i))
(48)
where the last step is due j < and j * pa(i) leading to the distance of nodes d(j, pi)) 2
In the proof, we adopt the following two lemmas from Nichani et al. (2024):
Lemma (Lemma 17 in Nichani al. (2024)). Suppose the transition kernel α of a Markov chain satisfies max,ey π (T' (s) > for all s' E V. Then, for any f-divergence D, and probability vectors x. y. we have:
Dy(ı o x o v) a (z(2ll1).
where the contraction coefficient is defined as:
== max (V) 57k
= max ||-j- - π (-11k]11-
2 jæk
Lemma 6 (Lemma 15 in Nichani et al. (2024)). Suppose ming, T(ss') z - Then we have
max (|+5-) - * <-11a)112 < 7
2 J
19

---

Published as a conference paper at ICLR 2026
C.5 PROOF OF LEMMA 4
Proof. Note that in the target inequality, the LHS equals - H (Xx Xpa(h)), while the RHS differs from but can be transformed to H (x, (xy). This result holds for both discrete and continuous x. For p-) and q(-) which are two distribution, by KL divergence's non-negativity, we have:
| p(s) log (()) z / log
Hence can get:
RHS
P(xh- 3, Xp = Pog s s/xp)(h)) =s))
/ P(XN =s') / |P(xx s|xx s') log IP(x) s s|xpa(h) = =')
(49)
$ -- J. P((N = (') = - (x(xn
where H(xn|xx>) (xx)) - T(xx; defined Definition By Lemma 3 and T(xa; Xpo(k)) >
have:
RHS & [(x)xx)) (xx)) < [(xx; xpo(s)) H(x)) L.S.
(50)
C.6 PROOF OF THEOREM 2
Proof. Note that the conclusion extend the dynamical system setting. 3 For notational gener- ality, we denote the transition kernel by p(- |-) in the following proof; this corresponds to π-- |-) in discrete Markov chain setting.
Recall that the transformer and BMA have the formula in Eq. (2):
exp(E, Ei log p(ar,,,(r))))))
AK Σ [1 exp 21 log (2)))))))))) Ene "n"-"""
(51)
where we define U/"-AAr exp(E, By the law of large numbers, have:
log log p(z (x(,0) E, log p(z),, x//)).
lim Σ log Bogr(ar (=n))) = ((oog((aaaa))) L
p(an) (2 pm()a) = lim
Σ log (a
L
Let Ih.N' A
Ei logp(a/ ). For all h" + pa(h), we have:
Un"" +po(h) = exp Σ log Logp(zhhahi ( Σ log p(ah w(aam)) = exp (L(ân,n" - -Dmpp((h))) +
L + 0o and lim (ōn, (k))) < O. Hence, we have limg An-ppa(h) = 1.
C.7 PROOF OF THEOREM 3
Proof. First, the transformer as constructed can be simplified as:
f (simp) (- H) = π 2i
Σ 1T i:hh Way,) E Rdd
(52)
Considering p E, at s1-h1 W // when W then p o(p) 1 and:
fo. (- H) π j=( where (( h
EN--
(53)
h'E[h-1]
In this case, the strong DPI follows from Lemma by verifying that the equality condition does not hold via direct calculation of covariance among Gaussian variables.
20

---

Published as a conference paper at ICLR 2026
Then based (p( diag(p) pp' op
computing the gradient of W w.r.t loss é in Eq. (12) yields:
ae(0; h
Ex 1(
2/h
dp
(x))
+
Öp
= Exx 1(
32h.
T
1
(IT
21:h-1 - π p(a - 1)
foo
(2h)
+
h -11
Eg. (53)
1
Ex( [(
( (TTTa1:h11 01:h-11 - (2 )
1
foo (10h) +
1
TT( (2 (x1)
T ( [20 (20h-1))
Ex
1/--1
1
Sou (2n)) +
So (Ir)) +
h Ex
1
T ( (x1)
(x (0n (21
1 RR-11
10 (ax))
1r (2:n)
where e and (s) = 1= (s) for all E V by assumption.
Then let ghh denote h'-th entry in a610o:h.9) Up
Rh-1 (N' [h - 1)), we have:
πT (atk xwwe) 1 (xn)
(ssss))((z = IN 88) -1)
Ex
1]
(54)
By Cauchy-Schwarte Inequality and Data Processing Inequality, we have:
i <x>)(<)) 1] = (slss))((nn = s, Ts = s') [
Ex]
(*x)))
1r
(55)
1 $ 5 (1==(100 aw) + Tx= pa(t))) S (arni ppo(k)) Ex
IT (x ppp(hh) 1M (1(a)))))
Eq. (55) has shown the desired result 2 S (pp((k)"
Assumption (Assumptions on transition kernel (Nichani et al. (2024), Assumption 1)), Ler 1 denote the spectral gap of n. We assume there exists
such that the following hold for
(Transition lower bounded): ming,s' ITs' {) > 7/M.
(Non-degeneracy of chain): EvEv [| s) 118 (-)11 2 /2/
C.8 PROOF OF PROPOSITION 2
Proof. In the dynamical system setting, the transition P(-I.) is given by the pdf of xı xn/xpa(h))
1 ||
pl y) (2 (1/2 22) exp
AA y
A E (RR))
2(1 p2)
Then with climinating constant terms in log p(v)/<w)))) with respect the candidate index h', get the equivalent form in BMA:
log p(ap (xr)
P
2i A Tt
Dy AAT my + const(h),
1-p
2(1 p2)
(56)
p2
pa N Ao 2(1 - p2) 11); PP(pa(N) 1:L = (p²)).
E=
Eq. (56) gives the BMA logits in the DS setting in sofumax form. We demonstrate that transformers under the observation restriction Eq. (7) cannot represent BMA in this setting.
). If we remove the assumption fo = j === ,²,,, Lemma 24 in Nichani et al. (2024) shows Ex LAx (wx)) 1] Ex (im,, mn,,)
*1eet fx = π = for (..).
-1] ATu where Tace is sequence length h divided by numbers of leaves of tree mish- Under Assumption and strong data processing inequality in Nichani al.
n(m, 1n
(*x))
(2024) (Lemma 5), we can obtain the non-asymptotic result (fa(M) 2h. 2 --102 20'
21

---

Published as a conference paper at ICLR 2026
Recall that, under Eq. (7), the transformer logits
4
P
Σ
Wer
2/1
1=1
while the BMA logits are:
P Cj
IT Aa
d
|
p² where CT = -T-p= 00, d 2(1-p") 0. Suppose, for contradiction, that the transformer exactly represents BMA, i.e.,
o (p a (péma) for all DS samples and all h E [M].
Since softmax invariant under adding a constant independent of h,, this means that for each fixed h there exists scalar b b(h) such that:
ple,
PENAN for all h E [h 1].
Using the DS model 21,, = PA x pa(h)) + V1 - p2 m/h we expand the logits as:
PEr, = P Σ Th
WorA's
a ppp(()) + VI p²
Ih Werth
I IT AA
L x pa(h)) + d Σ 1lat + ei +avi- p² Σ x Anh
cb PaHA,
c1P
1=1
Conditioning on all variables except Mer both sides of (+) become affine functions of the Gaussian noises n/. Since the DS distribution has full support and (*) is assumed to hold for all DS samples, the coefficients of the linear terms in { } must match for all realizations. Since (n/, J are independently sampled, comparing the coefficients of the linear noise term yields:
J IT (Wer e₁ A) = 0, Vnl e Rd + o (Wer - A) = 0.
Since in the DS model each a/,, is non-degenerate with full support, this forces:
W A.A.
Substituting W = A back into (+), the representation equation simplifies to:
L Σ
L
6+ C1P
The AA a paa(h)
= CiP
Σ
Th 1=1
IT
AA
apa(k) + d
|
1=1
Hence,
b d
|1 for all H E [h 1].
However, for a DS sample the quantities E, | 2 vary across and across samples, while is constant (depending only h). The only way the above equality can hold for all h and all DS samples is to have = = 0, which contradicts the assumption d * 0 in the BMA logits.
We conclude that no Wer make the transformer logits represent the B.A logits for all datasets generated from DS. Therefore, under Eq. (7), transformers cannot represent BMA the DS setting.
22

---

Published as a conference paper at ICLR 2026
C.9 PROOF OF CORRECTNESS OF CROSS-EXAMPES PARENT SELECTION
Lemma 7. Let μtt denote the stationary distribution. For any valid distribution μi A have:
(s) 88) S (sT(s') 8') log 8'), Vin AA.
8,8
Proof. Recall that for the transition kernel x 1-) and stationary distribution y=", it holds
1r (8) = Σ (8) (s) for all 8 E V. 2'EV
Fix an arbitrary jur E Ad, and for brevity write μ = 1=" We first upper bound the left-hand side. For any s' E V, consider the KL divergence:
KL(H || (1- (6)) = Ex log (F((s)
2
(T(s ((s(s') 5*)
sV
Expanding the inequality KL x(- (0))) 2 0 yields:
Σ
M n (s) Ts (8) s (s
C,
where the right-hand side C' does depend on s'. Multiplying both sides by pe(ss) and summing over s', we obtain:
Σ (3)))(() <
Σ p((s'C C- Σ
(57)
3.8'EV
S'EV
sEV
This bound holds for any choice of t4, E Ad. Next, we lower bound the right-hand side. For each s' E V, consider the reverse KL divergence:
T( s')
(L(T(:1 1()()))))))
Σ T(s 8') log SEV
2 0.
11a
Hence,
Σ i(s lo'g 2 Σ T(s (o)) (s).
SEV
SEV
Multiplying by n² (s") and summing over s' yields:
Σ (0((())( log (s (8) 2 Σ n (s')r((s s') nr (s)
8.E'VV
3.VEV
Σ 1 (s')T(() (i)ssss)))))))())) log 4 (s)
SEV
Σ
n" (s) log (s) =
(58)
sEV
where we used (x) = E. nt (F)((s)) (').
Combining (57) and (58), we conclude that, for any ias E Ad
Σ Pe(s')p (sg 10) SCs Σ (s()))((s) ( so) ((s s').
S
s.s'ev
23

---

Published as a conference paper at ICLR 2026
C.10 DEFINITION OF MUTUAL INFORMATION
Definition (Mutual Information and Conditional Entropy). Consider x, as two random variables in discrete continuous space Q. Lef xxy denote the joint distribution, and Pxx Py represent the marginal distributions. The mutual information I(x:y), entropy (x), and the conditional entropy H (xly) are given by:
[(x;)) [..[Pas Px log (xy(
H(x) - (2(
(x(
(59)
H(xly)= = [ log
P Pxy(2,)) H(() I(x;y), Py(1)
Further, x2 -muttal information is given by: (xx(xy)) = Ss (e,(,,)²²
F. 1.
D EXPERIMENT DETAILS
All experiments follow the same training setup unless otherwise specified: sequences are generated from a Markov chain with transition kernel x(- 8) ~ Dirichlet(a 14) with a = 0.1. We use batch size of 1024 for training and evaluate on 4096
samples. Parameters are optimized with Adam (Kingma& & Ba, 2015), using learning rate of 0.05 for discrete Markov chains and 0.001 for dynamical systems. For gradient- based analysis, we adopt SGD with learning rate 1. Fresh data are sampled at each iteration, and all implementations are based on JAX.
E ADDITIONAL EXPERIMENT RESULTS ON MARKOV CHAIN
Creas Extrapy Losa
60
Figure 8: Parent selection Lpa comparison between transformers and BMA during training. The metric is introduced in Eq. (10). The training configuration is the same as in the experiment shown Fig. 2.
Jmz
N,
36
14
Figure 9: Parameter-level comparison between transformer and BMA (W log r). Trainable W and Woy. Trained with = 20, H = 50, I. = 3, and 1024 training steps.
24

---

Published as a conference paper at ICLR 2026
(d-30)
0.7
3.8
0.6
0.6
05
0.25
15
04
0.20
03
0.15
02
40
0.10
0.3
6.05
10
20
25
Column tnoc
Column Indox
Column Ind
(xx (d=10)
Rrallog m) (d=30)
Grallog m) (d=550)
3.8
0.7
0.6
Bow Indee
0.5
Row
15
0.4
0.20
0.3
6.15
8-
0.2
0.3
0.
40
0.05
0.10
0.0
0.0
10
20
Column tnock
Column Indox
Column Indix
(Wc)(W)) -callog
[ear/Wel dax(log ml)
3.8
5-
2-
0.6
Box Index
20
15
0.4
0.20
6.15
8-
02
25
02
40
0.05
6.10
0.0
0.0
0.00
10
20
20
Calamn
Inder
Cakmm Inclev
Figure 10: Parameter-leva comparison of transformers with different vocabulary sizes where W.r has size E {10, 30, 50}. We can see Goal (Wr)) aligns well with aoal (log π) for all sizes d.
Ground Truth
Gradient Values ac/ap
Attention Predictions
Gooden: Stapp
-G
10
Figure II: Gradient Validation of op'
8f From left to right: ground-truth graph G, the gradients of
E l stacked row vectors, and attention weights A(2) (2) uniformly distributed since W 0.
25

---

Published as a conference paper at ICLR 2026
F
EXPERIMENT RESULTS ON DYNAMICAL SYSTEM
Head
Head 3
Head 5
2.0
10
100
-0.8
100
-88
103
-.8
0.6
0.6
0.5
200
200
402
402
403
250
Key
Head
Head 9
Head 11
3.0
1.0
200
-0.6
2.6
-06
200
200
203
0.4
5.4
0.4
300
300
100
200
100
420
300
200
300
400
100
200
300
400
Key Pos dons
Key Po5 Sons
Key Pastioons
Head
Head
Head
-066
200
200
203
300
300
303
0.2
5.2
0.2
0.0
0.0
102
420
102
220
4:0
200
22.0
300
Key Fox tome
Kay Fox tome
Key Pos
Figure 12: Visualization of Ist-layer attention (1) E RXT For readability, we visualize only nine of the twenty heads, better highlight the attention patterns on this long sequence length 400. The first layer replicates the historical occurrence of the same token. The model was trained with == examples, trajectory length H 20, vocabulary size d = 10, heads in the first layer, and 2048 training steps. The RPE parameters are initialized with a small positive value (0.5) along the construction direction, and grow to much larger magnitudes after training.
26

---

Published as a conference paper at ICLR 2026
dexample
344
24
53
344
Figure 13: 2nd-layer attention (4 HxH visualization. The attention patterns match the ground truth causal structure in the dynamical system setting.
Exoss-Example
Inner-Example RPE
-15
Figure 14: Visualization of the first RPE layer. The parameters are consistent with the construction.
W. KQ Matrix
W. OV Matrix
rt. [TTrnnstion Matrix)
-15
Courm 31der
Chlire
Calare Inlos
Figure 15: Visualization of the second attention layer. WKQ& pillxdl shows noticcable non-zero blocks on its diagonal. The occurring block is of size x d.
27

---

Published as a conference paper at ICLR 2026
Perent Selection Perfermance va Boyaston Model Averaghng
Crean Sro0p7 LaBa
1000
100
Figure 16: Parent selection loss during training in the dynamical system setting.
(Dal: feadind
aettern LAsn Curves
Pereet taleriee Less Teatyny vi
Offfrcare From BHL boodes Betrey
Figure 17: Generalization of parent loss (Ch') for transformers trained with I. E (1,...,30) ,003 in the dynamical system setting. Trained with d 10, H = 15, and 2048 training steps.
28

---

Published as a conference paper at ICLR 2026
G DISENTANGLED TRANSFORMER WITH ABSOLUTE POSITIONAL EMBEDDING
The disentangled transformer with absolute positional embedding (APE) is formulated as follows: Embedding Layer:
h [F(ww)) Pos( ( ex h
h
[0, e0
TR
1st Attention (K-head):
Attn. F: H (0 = hh ( W
(1),k KQ
(0).)
h1:t-1
(U)T
W
((2))
OV
(1),b
R
Disentangled Residual:
h2 = [/i (0) Attn; (H (0) , (H (0) ;0]
Attn
R
2nd Attention (1-head):
ful. H4) = (1 W
Wkgh KQ k())
h
(I)T w
Rd
1:t-1
(60)
First, we can see the model parameter W R.dxxdo where dg d T and T is the sequence KQ length. The total number of parameters in the first layer is O(d + 2²) compared to O H + L) parameters of the model with RPE Eq. (5). The redundancy of parameters may lead difficulties of interpreting the mechanism of transformers. Besides, for the disentangled transformer with APE, the embedding dimension is proportional to the length of input sequence, making it challenging to interpret transformers' mechanism on longer sequence tasks.
As for this transformer, we first provide theoretical construction which consistent with our construction for RPE model in Theorem 1. Empirically, we show this transformer can successfully select causal tokens. Besides, we provide results of trainable transformers showing alignments with construction in attention visualization and parameter verification.
G.1 THEORETICAL CONSTRUCTION
In this section, we provide a construction demonstrating how the proposed two-layer architecture possesses the capacity to implement the specific causal selection mechanism derived in analysis. Let the input embedding dimension be dg = d+T, where d the token dimension, T is the sequence length (due absolute positional embedding) and an input sequence contains L + 1 examples of length-L chain = HLL + 1). Suppose N1+1 denotes the set nodes from the last example, i.e., have NL+1 (t T 3h [H] HL + h).
G.1.1 LAYER 1: HULTI-HEAD ATTENTION CONSTRUCTION
The first layer consists of K attention heads (K < L). The Query-Key matrix (11,), specific predecessor tokens based on position. We construct it as a block matrix where the active interaction terms are confined to the positionallemmhedding subspace:
KQ attends to
0нжH
Oaxd
OaxT W (1). KQ
(1).k KQ
(1).
OT
Tн :
k-th block active), (61)
Orxd
0нжи
From this construction, if 3 + co, the attention weight of the first attention layer is given by:
if N
(1).k A
(62)
if NL+1, i= LH + h.
For the value projection, (11,..
propagates the semantic content of the attended tokens:
OV
(1(,),k OV
=
0xxd
Laxd
E (d+T)xd
(63)
'Following Von Oswald et al. (2023); Zhang et al. (2024), content embeddings (b(w)) of queries are zerved to prevent information leakage and seff-observation, as they are the targets of prediction. This strict sep- aration of input and target information consistently applied to the subsequent two Transformer architectures.
29

---

Published as a conference paper at ICLR 2026
And the output of the first attention layer is:
if 4 NL+1, A(011-1):
Attn. H 1
(L), (0)) A BIT w
(64)
1/ if E NL+1, LH +h.
G.1.2 DISENTANGLED RESIDUAL STREAM
Unlike standard summation of residuals, the disentangled transformer employs concatenation strat- egy. Nichani et al. (2024) proved that this transformer actually equivalent to a becoderrbbssd attention-only transformer (Theorem 3). The output of the first layer is the concatenation of the original input and the outputs all K heads:
Atin;
Attn
IRdo+Ka Kd
(65)
The dimension of the second layer input is dj = do + Kd = d ++T+ Kd.
G.1.3 LAYER SINGLE HINGE--HEAD ATTENTION CONSTRUCTION
The second layer employs a single attention head to aggregate the evidence collected by the K heads in the previous layer:
log π
Oaxa
Oaxd
Oaxd
Oaxt
OaxKa
Oaxd
log
Oaxd
W (2)
Oxxd
OXxT
OTxKd
(2 KQ
(66)
KQ
0 Oddxd 0 OKdxT
KQ
Oaxd
Odxd
log
Finally, the output projection W (
ov projects the aggregated context back to the semantic space by:
log π
W
Orxd
[R
(67)
Oxdxd
From the Eq. (64) and (65), we can see that h = [/, (0) xh...., x if E W = LH + h,
otherwise h [hg (0) .r,..., Ar]. Then, the attention score of the second layer for any N i == LH + ht of our interests, is given by:
Σ k- log (af))))))
for E NL+1, j LH + R'
(68)
Σx mJj log o
for j & NN+1-
So for i,j A2+1, have 1 EbE[K] log 7 (x ( aligned with Theorem 1.
Furthermore, suppose K L, i.e.,
L examples to infer the causal structure, and the Markov chain is stationary 27h 11. As L. - oo, for any j Nr+1, we have A /T. - E.,o (()))) log T(sls')
(s) (T). While for the true parent token HL + pa(h), have L E.,' n= (s')(((s)))) (s(s'),,,, which larger than (s) (s) The above quantity relation is drawn by the non-negativity of KL divergence, whose detailed proof is provided in Appendix C.9. Then, the attention weights of the second layer can select the causal parent token: (h) E arg max, lim Aij. And Vov predicts the transition.
Empirical Verification. We show the parameter visualization of the construction in Fig. 18a and its empirical attention visualization of parent selection in Fig. 18b. In Fig. 19, the constructed model shows precise parent selection accuracy (cross-entrops loss 0.0706), which is very close to the target algorithm BMA's (cross-entropy loss 0.0473).
30

---

Published as a conference paper at ICLR 2026
Will
10
Wyo
o(W.c.)
Top S x block of otWel
ITransition Matrix)
-0.1
-2
(a) The parameter visualization of theoretical construction.
Layer Selection
Multi- Attention Layer
Cuary Pastira
Medell Aiteation the 11
Madel Actuntion
Medel Atterdion #1
Famer
Parmi Cardilittrr
Tey Peallore
Ney Poollorr
Key Toatore
Kay Pealtore
Noy Pooldor
Kay Nortore
Fans:: Sodes pal1
Nearn wwdecccapppii palto
Panert Habis pa7t
(b) Attention pattern visualization of theoretical construction. In the first layer (right), the previous L examples are copied the hidden space of the last example. L In the second layer (left), the attention weights attend to the correct causal parents, which located the last 5 columns. The queries do not attend the keys from first I. examples empirically.
Figure 18: Parameter visualization and attention pattern visualization of theoretical construction.
31

---

Published as a conference paper at ICLR 2026
G.2 EXPERIMENTS OF TRAINABLE TRANSFORMERS
In the following, we train the standard disentangled transformer formulated by Eq), To show the alignment with theoretical interpretation, we use three strategies initialize the network: (a) fully random initialization: all the parameters are initialized randomly with Gaussian distribution; (b) block-amppp aaplliiii random initialization: parameters are initialized randomly (of scale 0.1), while the targeted block of the attention projection matrix is assigned a larger magnitude (of scale 0.5) introduce an inductive bias; (c) direction-consise
initialization: parameters are initialized such that the dominant blocks point in the analytically derived construction direction, still allowing model learning to refine the magnitudes (initial magnitudes: 0.2x optimal parameters).
We first compare the parent token prediction performance of these models during the training process Fig. 19. The results show that the 2-layer transformer is fully capable of selecting causal parents in its 2nd-layer attention head.
Then we visualize the attention pattern of the trained model in Fig. 21. For the first attention layer, the figure shows that queries from the last example L + mostly attend one example among the context examples, while some heads demonstrate degeneration with uniform attention to previous tokens. For the second attention layer, the transformers with different initializations all show their noticeable capability of predicting causal parents. Further, we visualize all the parameters of the transformer in Fig. 20. We can some alignments between the construction in Fig. 18a and the trained parameters. Since the transformer with absolute positional embedding has far more parameters of ((W&)) (1). KQ w w KQ' w than the one with RPE, the full interpretation of
its first layer diffieult. For the second layer, the parameter W (2) KQ also shows the diagonal pattern
consistent with construction and ( shows the log pattern.
Parent Selection Loss Comparison
DVA (pe-0 (20-0,2073)
Conctructed TF (0000)
Sandom int.
Bockrandom Int.
Consruct, dircctin
10.
300
Training
Figure 19: Parent selection loss Cpa of the transformer with absolute positional embedding and different initialization strategies.
32

---

Published as a conference paper at ICLR 2026
Thp x block o(Wed)
ITrandtion Manrid
dbb
(a) With Fully Random Initialization. Head 1 and 5 of the first layer / exhibit identity submatrix (5 x 5) at last column.
Head
Head
Head
Wis
o/Wcc)
Top S block of o'Woo)
(Trnnition Nortlo
(b) With Amplified Random Initialization. Head (1 in attention visualization Fig. 21b (Head 1). Head 5 shows multiple identity submatrices which possibly suggests superposition.
XQ degenerates which can he verified
Will
WT
WTl
WFll
o/W.oo
alW..)
Colarn
(c) With Direction Conssisentn Initialization. Head I and 5 of the first layer (()) submatrix the column which aligned with the theoretical construction.
KQ exhibit identity
Figure 20: Parameter visualization of trained transformer with absolute positional embedding. The second layer shows strong alignment in diagonal patterns WKQ and log pattern of Wov.
33

---

Published as a conference paper at ICLR 2026
Layer
Iead
Iead3
Hraad
Amerddnn
Kmy Proters
Greasd mur 2el I8C
Ground ires DAS licet
Sreane frmm its
fattors
No
(a) With Fully Random Initialization. In the first layer, Head 1, 3, and 8 of KQ matrices copy tokens from previous examples (to the query token the example), while Head 10 degenerates showing uniform attention (uniform features are seen as constants eliminated by 2nd softmax attention layer).
2nd Attention Parent Selection
1st Multi-lead Attention
Fonert 1008 Piro
P part
(b) With Block- Amplified Random Initialization. In the first layer, Head 3, 5, and 10 of KQ matrices copy tokens from previous examples, while Head degenerates showing uniform attention.
Attention Layer Purent Selection
Layer
Hend
Hend
Macid
Cardicadi
Key Featood
0.
Focton
No Furtou
Greand sith 246 (is
Ground trus DAS liccd
Greend truh 1sS
Km
Partora
No
(c) With Direction-Consisee Initialization. In the first layer, Head 1, 5, and 10 of KQ matrices copy tokens from previous examples, while Head degenerates showing uniform attention.
Figure 21: Attention pattern visualization of trained transformer with absolute positional embed- ding.
34

---

Published as a conference paper at ICLR 2026
H DISENTANGLED TRANSFORMER WITH APE VARIANT
H.1 MODEL ARCHITECTURE
For the variant of types of APE, we consider the disentangled transformer simplified by elim- inating some components added to residual stream. The transformer structure we consider below can be seen as substituing the positional embedding of sturcture Eq. (60) and simplify the model by assuming zero blocks in model weights:
Embedding Layer:
h2 (0)
= Pos (w)). (w)h
E RR
1st Attention (K-head):
Attn. oh 1
W (.).kk (1)) T #1=-1 W. (2).
KQ h
RR
Disentangled Residual:
h (1) !Attm!...,,Aiir |
IR Kd (69)
2nd Attention (1-head):
(1)T fee H4) ( Wk hn) 2 - I
R
where simply assume (1),& Ia- Since the difference lies in the positional embedding, the construction in Appendix G remains valid which can exhibit capabilities in causal token selection empirically. Besides, train this transformer under the same Markov chain setup in the trans- former with RPE experiments, obtaining consistent results as shown below.
H.2
EXPERIMENT RESULTS
Heads
Head2
Read
Headd
Hoad
Figure 22: Ist-Layer Attention Visualization of transformers in Eq. (69). Heads 1, 2, 5, 6 and
exhibits the diagonal block pattern the last rows performing the copying mechanism, while Heads 4, 7 and 10 degenerate to uniform attention. Heads 3 and give uniform outputs not influencing the 2nd attention layer (eliminated by softmax attention). Trained with 5, = 10, d = 5 and 10000 training steps.
35

---

Published as a conference paper at ICLR 2026
2nd Leper Athendion (Crumple 13
2nd Layer Attention iExemple 2)
2nd Laper Attention (temppe 3)
in
Truch (Exemple
DAG
Trush (Beample
204
94
L+l Figure 23: Visualization of 2nd-attention layer. Queries from the last example #1H Keys 1:L F1.T 2 1:H1 the whole sequence. Attention layer of disentangled transformer recognize the causal structure in-context.
wil Had
Head
Heac
WL Held
wiS Held
Wid Head
Figure 24: Parameter visualization of the first attention layer (1).k (10 heads in total). Full inter-
pretation is still challenging for huge parameter space. The attention-level behavior understanding can be referred to Fig. 22.
Weo
(ay)
(Transition Matris)
001
A56
A25
830
-256
425
Caland
KQ ov- In the variant of two types of absolute positional embedding, the second layer also shows strong alignment in diagonal patterns of WKQ and log ा pattern of Woy.
Figure 25: Parameter visualization of the second attention layer W 2))
W
36

---

Published as a conference paper at ICLR 2026
1 STANDARD TRANSFORMER WITH FEEDFORWARD NEURAL NETWORK
In this section, we consider a standard 2-layer transformer with FFN layers as follows:
Learnable Embedding:"
h
(0) (anby(w;) + Embp(w:) (w.),
MHA Layer & Residual:
h h1 MHA, (H W WRQ Wov),
E Rd
(70)
FEN Layer & Residual:
h ((+1) h + FFN, (1) W. b)
E R
Unembedding Layer:
Ser(- He)) = W (L)
Rd
1. the multi-head attention (MHA) formulated by: where HI))) [h]
MHA, H " 0) Σe
1DT (0 KQ k!0)
T ( (0) h1a-1 Wov
(71)
and the FFN layer formulated by:
FFN, H (00) W2 ReLU(W,,,, + b1) b2-
(72)
We consider the two-layer transformer L 2 with K heads in the first layer and one head in the
second.? For the task, the input sequence consists of .M 10 in-context examples of Length- H
Markov chains with d = 5 states and the total length T = H(M + +1). We set the hidden dimension as d' 128. For initialization, the parameters W of the transformer are initialized randomly by Gaussian initialization: Wij NN0 1/dw) where dw is decided by the dimension of W. We optimize the model using Adam with learning rate of 1 x 10-3 and a weight decay of x 10 Fresh data are sampled at each iteration of training without repetition.
Experiment Results. We train two transformers with 5000 steps and K 5 or 10 heads in the first layer. We observe the attention weights of the first layer visualized in Fig. 26a and Fig. 27a implement the copying mechanism where the features of one context example are copied to the position of the last example M+ the heads of the first layer show a diagonal submatrix occurring at the last several rows of example M + 1. Except for these, the remaining primarily show degenerated attention patterns at the rows of the last example M + 1. In the visualization of the second layer, we find that the trained standard transformer with MLPs can recognize the causal parents in its attention weights of the second layer. The aligned attention pattern and graph ground truth in Fig. 26b and Fig. 27b support our construction of how transformers can handle in-context causal learning.
Quantitative Results. We provide the results regarding how accurately transformers during train- ing can select random parents in their second attention layer in Fig. 28. We use the cross-entropy loss as the evaluation metric for accuracy and compare the trained transformers with BMA. We observe that during the training process, standard transformers gradually acquire the capability of in-context causal learning and approximate the loss of BMA.
SSimilar to the setup in Eq. (60), the query content embeddings are zero-assigned to prevent information leaking.
"Our implementation is based the codebase provided by Nichani al. (2024).
37

---

Published as a conference paper at ICLR 2026
(a) Visualization the first multi-head attention layer. Heads 1, and show the diagonal block at the rows of the last example. Information from previous examples is copied to the hidden space of the last example.
2nd Laper attandio (saample 13
2nd Layar Attortion licample 2)
2nd Layer Attuntion Sisample
23
23
M
Grousd Truck BAC (esemple 3)
Ground Trath DAG (tsampte 21
Grouse Trisk BAC (Bxample 3)
04
Fecnl Boder 8024
(b) Visualization of the second attention layer. Queries are from the last example #1:H the standard transformer recognize the causal structure in context.
2-+-1 The attention layer of
Figure 26: Attention visualization of the standard transformer with MLPs (5 heads in the first layer).
38

---

Published as a conference paper at ICLR 2026
(a) Visualization the first multi-head attention layer. Heads 2, and show the diagonal block at the rows of the last example. Information from previous examples is copied to the hidden space of the last example.
Pur Doy Proton))
Pure Bayy Priton))
Ground Truch Oicample
Groand Das (kample
Ground Truch (icample
L++1 The attention layer of (b) Visualization of the second attention layer. Queries from the last example æ1.H the standard transformer can recognize the causal structure context.
Figure 27: Attention visualization of the standard transformer with MLPs (10 heads in the first layer).
Parent Selection Loss Comparison
16
0.0473) Transformer w/ heads
Transformet w/ 10 heads
12-
Parent Selection
10-
1000
2000
3000
4000
5000
Training Steps
Figure 28: Parent selection loss Cpo of the standard transformer with learnable positional embedding and MLP or 10 heads in the first layer). During training, standard transformers gradually acquire the capability of in-context causal learning and approximate the loss of BMA.
39

---

Published as a conference paper at ICLR 2026
J PARAMETER-LEVEL VERIFICATION WITHOUT SOFTMAX
In the main experiments, we have compared the discrepancy between Gewil Wet and acn (log nT) Both of them have scale of [0, after softmax. Here directly compare Wir with log π + la without softmax. In the following, (Wer 1a is denoted as W for convenience.
To make a reasonable comparison between W (Wer 1aT) and log n, reference is needed to assess their difference since WW log Tayl [0. oo). Here, use log T|| as the reference and adopt the normalized RMSE (NRMSE):
W log 1
We log Tty
NRMSE( W, log T)
(73)
log T
log ||F
ij
We trained 3 models with vocabulary dimensions of {10, 30, 50). In Fig. 29, we first visualize the weieentwisee residual [Wi - log Tt =| (NRMSE, to preserve scale) and [Wa log mtyl2/11 log E12 (NMSE, to highlight structural deviations) for each model, and compare it with the original log From the 2nd and 3rd columns of the figure, the residuals are generally within small region, although some entries have noticeable differences. For these entries, the original log I itself has extremely large values in those place, i.c., Ttj 0 leads to large log Tij-
Quantitatively, we characterize the empirical distribution of each element W - log , // log ा || in the residual matrix. From the last column of Fig. 29, over 90% of the relative error between entries of W and log π lies within the region (0, 0.01).
Combined with heatmap visualization, only at the positions where Tay has values close to 0 (e.g., T2.6 e 40 log 22,6 -40 in Fig. 29, first column), the Wtr 1a has relatively large approx- imate error w.r.t. log π (over 0.01 but under 0.13 observed the last row of Fig. 29). Considering that log क (-o0, has unbalanced distribution while the trainable W initially has values uniformly around 0, this difficulty in learning small Tij may explain the large approximation error.
Further shown in Fig. 30, with continued training (up to 50K steps), the trainable Wer gradually increases its entry magnitudes, moving closer to the boundary values of log x. Overall, the results of Fig. 29 and Fig. 30 show that the trainable transformer parameter Wer approximates the BMA solution log π + la in the parameter space, as stated in Proposition 1.
Additionally, Fig. 31 shows that most entries of aeal (log are close zero. And after softmax operation eliminates the numerical differences from small values augmented by log (e.g.. e vs. 40) Toal (Wez) aligns with coa (log π)) across different steps.
Remark. The E Rd in the experiments shown in Fig. 29 and 30 is solved by the optimization problem arg mina Wtr la" - log π T||F. Since has closed-form solution, the column- mean of the residual matrix, we compute it directly.
K THE USE OF LARGE LANGUAGE MODELS (LLMs)
We acknowledge the use of LLMs like ChatGPT primarily to refine the grammar and improve the presentation of the paper. Besides, they are employed to polish the math proofs where the results validated by the authors. LLMs also assisted in writing portions of the experimental code, particularly for data visualization, which were reviewed and verified by the authors.
40

---

Published as a conference paper at ICLR 2026
log (d=10)
log n (d=30)
log (d=50)
-10
10
Row indox (1)
20
Bow mox 08
40
15
-60
25
40
-80
0)
malumnn
Normalizedd RMSE (dm 10)
Normalized RMSE (dm 300
Normalized RMSE (dm50)
0.34
5ow Indox 2
0.32
30
0.08
3:25
Now indor
0.10
0.06
15-
0.08
0.04
0.10
20
0.02
6
15
25
30
index ())
maliummn Index
Normalized MSE (d==10)
Normalized MSE (d=30)
Normalized MSE (d=50)
5.12
0.008
0.10
Row Rowernooo di
10
Now Irdeod
30
20
2.02
0.000
13
20
25
20
30
Co Columin index ()
Co kmt indixx 10)
Columin Indix (1
Aormalized RMSE ECDF (d=10)
Normalized RMSE ECDF (dm30)
Normalized RMSE ECDF (d=50)
1.0
2.0
0.8
EDF
0.6
0.00 0.02 0.04 0.06 0.08 0.10 0.12 (W- 1oo-loo--220 [8(ooorl)
2.000 0.005 0.015 0.015 0.020
0.000 3.002 0.004 0.006 0.008
Elements (W- Mggn-aa/M(kunn
Elements [I Villogml
Figure 29: Visualization of log R. element-wise NRMSE in Eq. (73), NMSE and the Empirical Cumulative Distribution Function (ECDF) of NRMSE. Larger approximation errors for log occur at entries where Tij is numerically close to zero. The ECDF plots show that large proportion of the entries approximate log with a small error, with over 90% falling within a relative error of 0.01. Trained with H = 50, L = 3, d E (10,30 30, 50}. and 2048 steps.
41

---

Published as a conference paper at ICLR 2026
T--3330
kgg
0.37
0.35
0.35
T-1020d
0.00
0.34
0.53
0.33
0.52
0.43
0.33
bicresing traiiing
(a) Visualization of Wtr and log
0.06
33000
0.000
0.322
0.34
6.44
0.52
0:42
0.86
0.000
6.00
0.00
76
82
15ke nowaing
(b) Visualization of W la" T and log x.
Figure 30: Histogram and density plot of Wer and W.r - 1a with increasing training steps T (10K, 50K). The entry magnitudes of Wer gradually decrease towards -o0. The density is estimated by kernel density estimation (KDE). Trained with H 50, = 3, d = 10, and 50K steps.
323000
7-20005
T- 5X00
5-
65-
1...A. w.T.T........ Iin....
Figure 31: Histogram and density plot of Gze (W.) with increasing training steps T (10K, 30K, 50K). Density is estimated by kernel density estimate (KDE). a (W:)) shows an aligned distribution with that of the task parameter aea (log x). Training setup the same as above.
42


Reviews:
[{"review_id": "3EjHjz0j0L", "reviewer": "Reviewer_EPSS", "summary": "This paper investigates how transformers can learn and adapt to different causal structures in-context. The authors propose a framework using Markov chains with randomly sampled parent dependencies and prove that a two-layer transformer with some version of Relative Positional Embeddings (RPE) can implement Bayesian Model Averaging (BMA), the optimal algorithm for this inference task. The work provides a detailed theoretical construction, information-theoretic guarantees, and empirical validation showing that trained transformers learn to approximate this mechanism.\n\nThe paper's contribution lies in the connection to BMA and its 2-layer RPE construction. However, the manuscript suffers from a significant weakness in the way in which it acknowledges and contextualizes these contributions with respect to relevant prior works. Additionally, the theoretical construction relies on several specific assumptions that tailor the architecture to the specific task.", "questions": "1.  **Unused Heads (Lines 294–295):** The authors state that \"some heads... didn’t learn meaningful features.\" According to the construction, each head in the first layer should learn a specific copying mechanism. Does this observation imply the model was trained with more heads than theoretically necessary, or that some heads failed to specialize as expected?\n2.  **Attention Notation (Lines 246–252):** The equations appear to use both a generic activation `σ(⋅)` and `softmax(⋅)` in the context of attention. Could the authors clarify the distinction or correct the notation if they are intended to be the same?\n3.  **BMA Formulation (Lines 138–140):** For clarity, it would be helpful to explicitly write out the analytical expression for the likelihood `P(x | pa(h))` and the resulting posterior-predictive (BMA) distribution. This would make the target algorithm that the transformer is shown to implement immediately clear to the reader.\n4.  **Theorem 3** Can the authors clarify under what distribution the expectations in Theorem 3 are taken, and how the result holds when the underlying Markov chain is non-stationary or non-mixing?\n5. Could the authors clarify whether a transformer with conventional architecture and positional encodings (without separating L position and H position) would still approximate BMA under the same task formulation, or if these design constraints are essential for the claimed behavior?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "1.  **Principled Connection to Bayesian Inference:** The paper provides a link between the transformer's attention mechanism and Bayesian Model Averaging (BMA). \n2.  **2-Layer RPE Construction:** The explicit construction of a 2-layer transformer using a two-axis RPE to implement BMA reinforces the idea that architectural components can be mapped to algorithmic steps.\n3.  **Information-Theoretic Guarantees:** The paper provides additional theoretical guarantees for causal structure recovery based on mutual information, strengthening the formal understanding of the learning dynamics.", "weaknesses": "1.  **Insufficient Acknowledgment of Prior Work:** The paper lacks a proper discussion of the related work  [1]. It does not properly contextualize its contributions, preventing a clear understanding of the paper's own contribution over prior works.\n    - **Framing of Novelty:** The introduction frames the manuscript’s contribution as a departure from prior work limited to \"fixed dependency structures\" yet [1] seems to have already moved beyond this by introducing the same setup considered here based on in-context causal structure selection in Markov chains. The section \"Our Approach\", which conceptually outlines a setup closely related to [1] (each token depends exactly on one of the past tokens, with this dependence inferred in-context), omits any mention of this direct precedent work.\n    - **Appendix Citation:** The citation in the appendix appears to report incorrect information. It claims that the verification in [1] is \"limited to attention visualizations\" however, [1] also provides quantitative validation using KL Divergence. It also criticizes [1] for \"task-specific constructions,\" a trait shared by this work (see points below).\n    - **Unacknowledged Theoretical Precedent:** the work presents in its Lemma 2, the central statistical argument proving the identifiability of the true causal parent. However, the result seems to closely parallel Lemma 2 in [1].  The present version is slightly stronger (the strict inequality) and presents a different proof, but leads to the same conclusion. This fact should be correctly acknowledged. \n    *   **Setups overlap** It would be helpful if the authors could acknowledge the conceptual overlap and similarities with [1]. The key difference seems to me to be that [1] infers a single global structure from within one sequence, while this work infers position-specific structures from across multiple sequences. This makes this work a valuable generalization and a genuine advance in the expressivity of the underlying graphs, but also requires the number of heads as well as the embedding to scale with the length of the examples (H). The two frameworks address the same underlying estimation problem: at a given position, identifying the correct parent accumulating log-likelihoods in a one-parent Markov process,  but differ in how the evidence is presented in the context (within-sequence vs. across-sequences). Acknowledging this would strengthen the paper by highlighting its novelty while giving proper credit.\n    \n2. **Dependence of the BMA result on architectural tailoring:** While the paper presents interesting ideas regarding linking in-context causal selection to Bayesian Model Averaging (BMA) in transformers. It remains unclear whether the demonstrated BMA implementation genuinely reflects what standard transformers learn and can implement or rather a consequence of the architectural choices, aligned with the task structure, made in the construction.\n\n3.  **Different attention domains across layers (T×T then H×H)** Layer 1 attends over the full concatenated sequence (shape (T\\times T) across examples×positions) while Layer 2 attends only across positions inside an example (shape (H\\times H)). This effectively changes the logical index set the model operates on between layers, i.e. the model is built so different layers operate on different axes, tailored to the task they need to solve. It remains unclear whether a standard attention mechanism can discover a similar pattern.\n\n4.  **Two-axis positional encoding (separate lookups for example-index and position-within-example)**\n  The construction uses distinct positional biases for the example axis and the within-example position axis. This is a RPE that is tailored to the task and explicitly factors the task axes; it is not the standard single-axis absolute or relative positional scheme used in many transformer models. It remains unclear if replacing the two-axis RPE with a single standard positional encoding (absolute or standard RPE) the same BMA mechanism emerges or if it can even be represented.\n\n5.  **First layer omits W_QK and W_OV.** The first attention layer is constructed to behave as a direct copier (fetching particular past tokens), which is commonly employed in the literature. However, the matrices acting over the semantics W_QK and W_OV seem to be omitted in the construction. It is acceptable for the purpose of the theoretical construction to fix these components (W_{QK} and W_{OV}​) to zero, but this design choice should be stated explicitly. Moreover, the paper should provide empirical evidence showing that allowing these parameters to be trainable does not substantially alter the main BMA mechanism or invalidate the proposed construction.\n\n6.  **No positional encoding  for the second attention layer**. Similarly, the construction does not handle positional information consistently across layers: the second attention layer either omits or fixes to zero positional encodings compared to the first. This asymmetry effectively hardcodes aspects of the intended computation and departs from the symmetry typical of standard multi-layer transformers. As in the previous point, it would be important to assess whether including full positional biases in the second layer would materially change the mechanisms the model implements.\n\n\n7. **Theorem 3** Theorem 3 leverages the same underlying statistical mechanism used in [2] (as acknowledged by the authors): it expresses the gradient signal in terms of a χ²–mutual information measure between each candidate parent and the child token, and then applies the Data Processing Inequality to show that this measure is maximized for the true parent. The contribution of Theorem 3 is, in my opinion, incremental: it adapts the same statistical idea to the specific transformer construction that exactly implements Bayesian Model Averaging. Moreover, the paper states that its proof “eliminates the stationary assumption of the data distribution and doesn’t require the Markov chain to be mixed,” but  Theorem 3 still relies on expectations such as $(\\mathbb{E}[\\pi(x_h \\mid x_{h'}) / \\mu(x_h)])$, which implicitly assume access to a stationary marginal distribution or i.i.d. sampling across examples. Without some ergodicity or mixing condition, it is unclear to me how these expectations are defined or estimated.\n\n[1] D'Angelo Francesco, Francesco Croce, and Nicolas Flammarion. \"Selective induction heads: How transformers select causal structures in context.\" In The Thirteenth International Conference on Learning Representations. 2025.\n\n[2] Nichani Eshaan, Alex Damian, and Jason D. Lee. \"How Transformers Learn Causal Structure with Gradient Descent.\" In Forty-first International Conference on Machine Learning. 2024.", "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": "dVanXRV5fp", "reviewer": "Reviewer_M4B8", "summary": "The paper studies how transformers can learn causal structure in context. It designs a task where for each input sequence a causal graph is sampled and used to generate Markov chains (with fixed kernel). Given $L$ examples in context, the model has to learn to predict the continuation of the $L+1$-th sample. The paper shows that there exist a construction of a 2-layer transformer with $L$ attention heads in the first layer which can implement BMA to solve the task, and this is (approximately) learnt by the trained models. Finally, it studies the extension of this task and model to continuous data, i.e. a dynamical system.", "questions": "See above.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- The paper studies a novel problem, i.e. how transformers can learn causal structures in-context, which continues an established line of work on understanding the internal functioning of transformers in simplified setups.\n\n- The paper studies the proposed task both theoretically and empirically, showing that transformers implement BMA, and even provides an extension to the continuous case.", "weaknesses": "- The paper contains a lot of different parts, which makes it cluttered and none of them is presented too clearly. For examples, the construction of the transformer to implement BMA is not actually given, the figures are too small to be readable, there's no conclusion section. Moreover, the notation is in several places imprecise, e.g. in L292 $x$ is discussed but doesn't appear in the previous equation, in Eq. (6) $k'$ is not defined, etc., and the writing at times unclear. Overall, this makes the paper hard to follow.\n\n- The construction of the transformer which implements BMA needs the number of heads to scale with the length input sequence, which I think it's a limitation and makes the construction even more distant from real-world models.\n\n- While the paper provides several results, it's not clear what the takeaway message is. In fact, the paper shows that transformers can implement BMA to learn a causal structure in context, but this is shown by tailoring the architecture to a specific task (see point above). This construction doesn't seem to reveal some general mechanism applicable to other tasks, as it was for example with induction heads, which may limit the impact of this result on future work.", "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": "1CStmE33JR", "reviewer": "Reviewer_h4T5", "summary": "This paper extends prior work on analyzing how transformers learn causal structure by proposing and analyzing a task which requires estimating and performing inference on an unknown causal graph from a number of in-context examples. The main theoretical result is a construction of a two layer disentangled transformer which can compute the posterior distribution over causal graphs (Theorem 1). It then uses experiments on disentangled transformers to support the claim that transformers learn a similar construction.", "questions": "- It seems that in addition to reducing the number of parameters, factoring the positional embeddings into L and H also simplifies the construction since it gives you parameter sharing between the in-context examples and the test-example for free. How do the experiments change if you use a more standard architecture?\n- Is it possible to interpret the learned attention heads and show that it really does implement the same construction as Theorem 1?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The toy task is interesting and captures the idea of adapting to an unknown causal structure\n- The paper is generally well written\n- The authors show consistency of the construction (Theorem 2)", "weaknesses": "- Theorems 1 and 3 aren't mathematical theorems due to the use of $\\approx$. For example for Theorem 1 I believe the intended theorem is something like there exists a sequence of weights $\\theta_\\beta$ such that $\\lim_{\\beta \\to \\infty} f_{\\theta_\\beta}(\\ldots) = \\pi(\\ldots)$? Similarly, the mathematical statement of Proposition 2 is also unclear to me.\n- As far as I can tell, the paper doesn't actually demonstrate weight-space agreement between their construction and the one learned by gradient descent (only the attention maps are verified). Is the challenge in matching the behaviors of the different heads? It's not clear at all to me from Figure 16 that the construction matches the one in Theorem 1.\n- It would be good to run some experiments on a standard transformer with MLPs to check whether you learn a similar construction or get similar performance.\n\nMinor points:\n- There is a redundant prove in the abstract: \"We **prove** that a two-layer transformer with relative position embeddings can **provably** implement\"\n- The figures in the appendix are very hard to read. For example, Figure 10 makes it look like every weight is identically 0 since the diagonal entries are barely visible. Perhaps putting the attention weights in log scale or reducing the size would help?\n- Typo on line 268: \"We set transformers has $K$...\"\n- I'm not sure what footnote 3 is saying?\n- Proposition 2 is cited at Lemma 2 in Appendix B.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": "Z4oCxYCctE", "reviewer": "Reviewer_9XhA", "summary": "This paper introduces a framework based on Markov Chains to explain how transformers create a causal dependency graph between tokens using ICL. In the Markov chain, each token depends on exactly one prior parent.\n1. The paper shows that a disentangled transformer can be trained to do Bayesian Model Averaging.\n2. Provides empirical evidence\n3. Extends the framework to continuous case (Linear dynamical systems)\n4. Provides a theoretical guarantee based on information theory for the transformer causal structure selection.", "questions": "Do the authors have any insight into how their work can be extended to transformers with MLPs and multiple layers?\nIf not, maybe an experimental section on standard transformers to back the claims empirically?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper makes a good connection between BMA and attention (for the disentangled transformer). \n2. Shows there exists a model which, by construction, implements BMA. Also supports the claim empirically. \nThe core strength of the paper is that it helps in formalising how ICL works through a probabilistic framework.", "weaknesses": "1. As mentioned in the strengths, the formalisation of ICL that the paper brings is useful; however, the main weakness is that both the proofs and the experimental work are limited to a special form of the transformer. It is unclear how the results can be applied to the standard architecture.\n\n2. Additionally, the paper claims that the proofs and experiments were conducted on a standard transformer. As far as I understand, they work on the disentangled transformer, which is not the same. Claims should be addressed to reflect the paper.\n\n3. It seems that the paper is missing a conclusion.", "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": "o790bITnDg", "reviewer": "Reviewer_BaHK", "summary": "This paper investigates whether transformers can learn causal structures underlying sequential data. Using a framework based on Markov chains with randomly sampled causal dependencies, it shows that a two-layer transformer with relative positional embeddings (RPE) can exactly implement Bayesian Model Averaging (BMA). It further provides information-theoretic guarantees explaining how transformers recover causal structures.", "questions": "1. Extension to nonlinear dynamical system:\nThe current analysis focuses on linear dynamical systems, while transformer is a nonlinear mapping.\nDo you expect that the this framework could be extended to the nonlinear dynamical system ? \n\n2. Essential role of RPE: The theoretical construction appears to rely on the use of RPE. Do you think other PE, such as absolute PE or learned PE, could, in principle, reproduce the same behaviors, or is RPE essential for causal structure inference?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper provides a rigorous theoretical construction showing that a two-layer transformer with RPE can provably implement BMA for causal structure inference. The information-theoretic analysis strengthens the understanding by proving that attention weights converge to the true parent structure under increasing in-context examples $L$, using mutual information arguments. So I think this paper make meaningful contribution.", "weaknesses": "While this paper is theoretically solid, some assumptions, such as linearity of the dynamics and  specific design of RPE, limit its generality. However, I think that these limitations are acceptable for the scope of this paper. Out of interest, I would like to ask the 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": {}}]

Author rebuttal:

Decision: Accept (Poster)