Transformers Learn Latent Mixture Models In-Context via Mirror Descent
Authors: Francesco D'Angelo
Nicolas Flammarion
Venue: ICLR 2026 Poster (2026)
Keywords: in-context learning
markov chain
transformers
mirror descent
mixture models
latent variables
Primary area: foundation or frontier models, including LLMs

Abstract: Sequence modelling requires determining which past tokens are causally relevant from the context and their importance: a process inherent to the attention layers in transformers, yet whose underlying learned mechanisms remain poorly understood. In this work, we formalize the task of estimating token importance as an in-context learning problem by introducing a framework based on Mixture of Transition Distributions, where a latent variable determines the influence of past tokens on the next. The distribution over this latent variable is parameterized by unobserved mixture weights that transformers must learn in-context. We demonstrate that transformers can implement Mirror Descent to learn these weights from the context. Specifically, we give an explicit construction of a three-layer transformer that exactly implements one step of Mirror Descent and prove that the resulting estimator is a first-order approximation of the Bayes-optimal predictor. Corroborating our construction and its learnability via gradient descent, we empirically show that transformers trained from scratch learn solutions consistent with our theory: their predictive distributions, attention patterns, and learned transition matrix closely match the construction, while deeper models achieve performance comparable to multi-step Mirror Descent.

Published as a conference paper at ICLR 2026
TRANSFORMERS LEARN LATENT MIXTURE MODELS IN-CONTEXT VIA MIRROR DESCENT
Francesco D' Angelo" TML Lab, EPFL Switzerland
Nicolas Flammarion TML Lab, EPFL Switzerland
ABSTRACT
Sequence modelling requires determining which past tokens are causally relevant from the context and their importance: process inherent to the attention layers in transformers, yet whose underlying learned mechanisms remain poorly under- stood. In this work, formalize the task of estimating token importance an in-context learning problem by introducing a framework based on Mixture Transition Distributions, where a latent variable determines the influence of past tokens on the next. The distribution over this latent variable is parameterized by unobserved mixture weights that transformers must learn in-context. We demon- strate that transformers can implement Mirror Descent to learn these weights from the context. Specifically, we give an explicit construction of a three-layer trans- former that exactly implements one step of Mirror Descent and prove that the resulting estimator a first-order approximation of the Bayes-optimal predictor. Corroborating our construction and its learnability via gradient descent, empir- ically show that transformers trained from scratch learn solutions consistent with theory: their predictive distributions, attention patterns, and learned transition matrix closely match the construction, while deeper models achieve performance comparable to multi-step Mirror Descent.
INTRODUCTION
In recent years, Machine Learning has been transformed by large language models (LLMs). These massive, complex models achieve unprecedented performance across diverse tasks beyond text gen- eration (Bubeck et al.) (003). A striking example is in-context learning (ICL) (Brown et al.) [2020] Min et al. 2022). where models adapt to new tasks using only examples the prompt without parameter updates. Mechanistic interpretability has made significant strides in explaining this phc- nomenon, revealing that transformers can implement computational circuits (Elhage et al.] 2021 Olsson et al.) 2022 D' D'Angelo et al. 2025) that mimic known algorithms. For instance,
set-
tings like linear regression, they learn to implement gradient- gradient-based optimization (Garg et al. 2022 Akyürek et al. 2022 Von Oswald et al.
2023a b. Zhang al. 2023 Ahn al. 2023 Mahankali
ct al. 2024), while for Markov Chains, they implement
coanting--bsedd estimators for the transition probabilities (Nichani al., 2024 Edelman al. 2024 Bietti al. 2023a Rajaraman al. 2024
Ildiz et al. |024| [Sete & Cotterell 2024] Chen et al. 2024). However, these successes are con- fined to problems where all the sequences rely on the same, fixed causal structure: the relationship
between tokens is static. For instance, in the case of regression, the model only needs to learn that every even token depends on the previous odd token and in Markov chains, learns that the next token depends only on the previous one.
Real-world sequential data, particularly language, defies such simplicity. The meaning of a sentence arises not from the fixed sequence of word meanings, but from the dynamic causal links between the words that must be inferred from the context. These underlying structures, which are fundamental to language, are latent variables hidden from direct observation. Consider the sentence in Figure I to predict the final token, a model cannot rely on simple recency. It must infer a latent structure: that "'og" is the agent, "ball" is the relevant object, and "bird" is distractor. The influence of past token is not merely a function of its position, but of its inferred role. This ability to infer and reason
"Correspondence Francesco. .aneeecoo.ddneeeee oBepf1.ch

---

Published as a conference paper at ICLR 2026
ferch chase
Low Ag
dog 0.6
0.3
0.0
0.0
High Xg
The dog that saw the bird threw the ball, and then ran to fetch. π High Au
ball
0.5
0.0
0.4
0.0
bird 0.0 0.1 0.0 ... 0.8
To predict the final word, the model must infer the causal relevance of past tokens. Our MTD framework models this by separating static, context-free unigram (x") from dynamie, context-dependent weights (A) that are inferred in-context. The model learns assign high weights the causally relevant positions ('dog') ) and ( ("ball", activating their respective slices of the unigram (e.g., æ (dog, -) favouring verbs like fetch or chase). Converscly, a low weight is assigned to the distractor ('bird"), ) suppressing its influence.
Figure 1:
unobserved variables, be it syntactic roles, speaker intent, or the causal topic of a discourse, is hallmark of intelligence. A robust model must therefore move beyond fixed dependencies and learn to infer this latent structure, dynamically identifying which tokens are causally relevant in given context. These considerations give rise to the core question of this paper:
Can transformers infer latent structures in-context, and what algorithm do they learn?
To investigate this question, introduce synthetic task based the Mixture of Transition Dis- tributions (MTD) model (Raftery) 1985, André Berchtold, 2002). This framework recasts token- importance estimation as learning latent mixture weights in-context. While prior work has used mixture models in regressions (Pathak et al. 2024) or HMMs (Xie et al. 2022) to study ICL, they did not unveil the mechanism through which transformers learn to infer the latent mixture weights. We make the following contributions:
1. We introduce framework for ICL based on latent variables using MTD: create syn- thetic task that frames the estimation of the influence of past tokens as learning latent mix- ture weights in-context.
2. We show that transformers can solve this task by implementing Mirror Descent and give an explicit construction of a 3-layer transformer that exactly implements one step of this algorithm.
3. We prove that the resulting one-step estimator yields a first-order approximation of the Bayes-optimal predictor.
4. We empirically validate that transformers trained from scratch with Adam learn solutions consistent with this construction matching the predictive distributions and attention patterns while deeper models achieve performance comparable to multi-step Mirror Descent.
We show that transformers implement Mirror Descent to dynamically infer which past tokens are relevant, a capability fundamental to sequence modeling and, in particular, to capturing higher- order dependencies in language. This provides gradient-based explanation of ICL in sequence modeling and offers a new algorithmic lens for understanding latent-variable inference in attention- based models. We defer to the Appendix C more detailed discussion of related work.
Why MTD? Our choice of the MTD model is motivated by the desire to capture both in-weight and in-context learning in a single, controlled setting. While prior ICL work has focused on settings and tasks where nearly all useful structure must be inferred from the prompt (e.g., gradient-descemt ICL in linear regression or counting estimators for simple Markov chains), our setup explicitly assumes that some statistical structure, namely the transition matrix T², can be stored in the model's weights during pretraining and reused inference time to infer in-context the mixture weights À from single sequence. This mirrors a plausible regime for large language models, which cannot memorize full high-order r-grams due to their sample complexity growing exponentially with n but can realistically encode lower-order 7-grams in their weights and dynamically reweight the influence of different past tokens depending on the context. Our in-context learning framework explicitly couples an in-weight component (learnnng π²) with an in-context component (inferring A). providing a natural testbed for studying the interplay between in weight and in in-context learning in L.Mss.
2

---

Published as a conference paper at ICLR 2026
2 DISENTANGLED TRANSFORMERS
The disentangled transformer (Friedman et al. 2023) a modification of a standard Transformer with relative positional encodings (RPE) (Shaw al. 2018). designed for enhanced interpretability by: (1) removing MLPs; (2) replacing residual connections with concatenation, creating an explicit residual stream that preserves computations from all previous layers; and (3) simplifying the atten- tion mechanism to use single attention matrix instead of separate query and key matrices, and absorbing the value projection into the output matrix. Nichani et al. (2024) demonstrated that disen- tangled transformers are equivalent to standard transformers using only attention layers. The model maps a sequence of tokens 8 = (s1:..., 8T) from a finite alphabet S a sequence of vectors. Each token s; is represented by its one-hot vector es, E {0, 11151, yielding the initial representation h ex; E IRdo with do |S]. The model consists of I layers. We denote by h(()) RRd the full hidden state of token i after layer L, and by h; (1.Ak the output of attention head h in layer I. For each (2,b) between query i and key j is:
head h E {1...., Hi), the pre-softmax attention score Cij
(1.h) (h. ) W h (1-1) + "
(.A))
(I)
Here, A E R4-1xx--1 is a learnable attention matrix and (1.) = (R (1.1)
rij
)-j++1; is relative
positional encoding vector retrieved from a learnable lookup table R (1.)) E Txd The attention !
weights are computed via a causally masked softmax: Al = |softmax (e, )1 where e ((LA) is
the vector of scores. The output of single attention head formed by concatenating the full hidden
state h (t-1) with positional value embedding ('(hhh) before the weighted sum. The full hidden state is then updated by concatenating the input with all head outputs:
Concat (n(1-1), (L.1)) ...., A A...M.).. (1.b) The positional value embeddings r (1.A) are retrieved from second lookup table R. E RxxRR, where dR is fixed hyperparameter. The dimension of the representation thus grows according to the recurrence di = d1-1 + H1 (d-1- + dR). After layers, a final linear layer with matrix Wo E x maps h to logit predictions.
h. (1,0)
(1.b) Concat (h (1-1) r (lllh)))
h
3 MIXTURE OF TRANSITION DISTRIBUTIONS FOR IN-CONTEXT LEARNING
The Mixture of Transition Distributions (MTD) model (Raftery 1985), is higher-order Markov chain that offers parsimonious representation of long-range dependencies. The idea is to model the probability of the next state as a mixture over lags, where a transition matrix is applied to the token each lag, and the mixture weights determine the relative influence of each past position.
Model Definition: Let Y = (Y....., .Y) be a sequence of random variables tak- (1....,q). The MTD model of order m ex-
ing values in finite alphabet y' plains this sequence by positing corresponding sequence of unobserved latent vari- ables
(Zm+1,.... Zr), where each Z, E [L.....m] acts as a switch, select- ing which of the previous tokens, Y--1...., Y will influence the current token Y,. The selection of this lag is a random event, governed by the mixture weights À = (A1,...,A)) Am) with Ag 2 for all g and Σg=1 X, =1 such that the probability of choosing lag g is given by g) Ag- Once the lag Z, = g is sampled, the next token Y, is generated from a first-order transition that depends only on the state at the sampled po- sition, Yz-g- This is captured by transition matrix
P with P the set of q x q row-stochastic matrices defining the conditional probabil- ities T [i, P(Y:
Figure 2: MTD for m=2.
Y--g i, Z. g). Marginalizing over the unobserved latent variable Zr, we obtain the predictive distribution:
Definition 1 (Mixture Transition Distribution). A sequence of random variables Y follows m-th order MTD model if for all > m and any history yi the conditional probability of Y, is:
P(Y = yt Y
yi >n
Xıg т (r(tt-g: V:),
(2)
g=1
where the parameters X (A1...., Am) are the mixture weights and the transition matrix.
3

---

Published as a conference paper at ICLR 2026
This structure allows one to model different effective contexts, making it more flexible than first- order Markov chain while preserving tractability by requiring only m - 1 + q 1) parameters compared to the qTth (q 1) of full m-th order Markov chain.
The In-Context Learning Task:
We design an in-context learning task based on the MTD model structured follows: Given transition matrix π we generate sequences y, by sampling for each new vector of mixture weights À from a prior X ~ Dirichlet( a = 1). This X vector is the hidden task that defines the statistical structure of that particular sequence. The sequence y = (.),...., yT) The learning objective for model F, such transformer, to predict the next token yt given the context y
is then generated according to the MTD model in equation 2
inf Birrcirichle((aa) MY-MMMAAA KL (p(Ylv) 1-1 X) f(y f()] fEr
To perform this prediction optimally, the model must effectively learn the mixture weights À of the latent variables Z. which are not directly observable in the sequence and differ across sequences:
In-Context Task: Learn the unknown mixture weights À given a single sequence y.
The Bayes Optimal Solution: The solution to the ICL. task (Eq. is the Bayesian predictive distribution, p(Y++1 o.)
Proposition 1 (Bayes-Optimal MTD Predictor). Given the MTD model with known transition matrix æ. a prior p()(e), and an observed sequence yi. the Bayesian predictive distribution for Y++1 = is a convex combination of the lag-specific transition probabilities, weighted by the poste- rior mean of the mixture weights X:
P( = v1,
i Bayes T (V6+1-992) 9
Bayer = E(>, ví cx) = [ A yi, ad dA, Jmmi
9
9=1
(4)
where p(A yi a) xs plvi | A)p(> α)) is the posterior distribution with plyf Il (Em= (y yk)) being the likelihood.
A)
The derivation is provided in Appendix F While elegant, the Bayes-optimal predictor is analytically intractable. The core issue is that the Dirichlet prior is not conjugate to the MTD likelihood, mean- ing the posterior distribution does not have closed form. Consequently, the integral defining the Bayes cannot be computed directly, necessitating the use of approximation methods. posterior mean Ag
Mirror Descent (MD): The intractability of the posterior mean motivates considering simpler point estimates, such as the Maximum Likelihood Estimate (MIH) the Maximum A Posteriori (MAP) estimate which are equivalent under uniform Dirichlet prior. However, analytical computation of either the MLE or MAP is intractable, necessitating iterative optimization methods. Methods such (EM) or Mirror Descent (MD) are preferred over standard gradient
as Expectation-Maxinue descent, as they are naturally adapted to the geometry of the simplex. For an extended discussion, see Appendix G Furthermore, even if the MAP estimate could be found, it represents the mode of the posterior distribution, which does not coincide with the posterior mean in the Bayes-optimal predictor, therefore leading to suboptimal predictions.
Mirror Descent, is a first-order method particularly well-suited for optimizing over the probability simplex A N & Yudin (1983); Beck & Teboulle (003). By using Bregman diver- gence based on the negative entropy potential (A) = (A), where (A) = - Σ Ag log X the Shannon entropy, MD results in the Exponentiated Gradient (EG) algorithm, which has a simple multiplicative update rule (see Appendix G.4 for details):
(k) exp(n (M))) )g)
1 (k+1)
A
(S)
E
A (k) exp(1) VAl(A()))))
where 7 0 is the learning rate and ()((A()) the gradient of the log-likelihood evaluated (k) Instead of iterating the EG algorithm convergence, we analyze a non-iterative es- timator derived from a single update step, initialized at the center of the probability simplex 1/m)). This approach yields a computationally efficient, regularized approx- imation of the MLE, which demonstrate serves an effective proxy for the posterior mean.
(A (0) =
4

---

Published as a conference paper at ICLR 2026
Proposition 2 (tne-Step MD Estimator). Initializing the EG algorithm (Eq. at A
(1/m,..., 1/m) and applying single update step for the MTD model yields the estimator:
exp (..m) Ek=m+t (k(9))
Tikk g. Vk.)
MD Ag
Tk(9) p(Zx = yr. y1. (0)
Ej=1 exp 1l Ex=m+1 (
ot (2k h: Vke
(6)
where (g) is the posterior responsibilityy of lag g at step m, under the uniform prior.
The derivation is provided in Appendix G.4
This one-step estimator computes, for each step k: in the sequence. the posterior probability that lag was responsible for generating yk (assuming all lags equally likely a priori). These responsibilities are summed across the sequence and fed into softmax function to produce the estimate ÂD with the learning rate v/ controlling its sharpness.
TRANSFORMERS IMPLEMENT ONE-STEP OF MIRROR DESCENT
We present our main theoretical result: a constructive proof that the one-step MD estimator can be implemented by Transformer. The crucial mechanism relies on relative position encodings to cor- rectly route information, allowing the self-attention layer to compute the posterior responsibilities.
Proposition 3 (Transformer Implementation the One-Step MD Estimator). Given an MTD model of order with a known transition matrix π P.for any sequence y1.r of length T 2 m, there exists three-layer disentangled Transformer 'T with single head, do = and dR 2 m. that implements the one-step MD estimator. The Transformer produces at position T the predictive distribution for token Yr given context y
T(yTT) =
Σ (B T (yT--::)
(y(hrr)) =
exp
T-m
3 DT
Ei=m+1
76(9))
(7)
E exp
9 Σ Ei=m+1 T-
with the weights \(yr-)) computed exactly as the one-step MD estimate and B3 is a learnable pa- rameter corresponding to the sealed learning rate of the MD algorithm.
In the following prove Proposition 3 by explicitly constructing 3-layer disentangled Trans- former that implements the one-step MDD estimator. The first layer computes the posterior responsi- bilities (g). the second layer computes the logits Ei=ma (g(g), and the third layer produces the final estimate vector \(y1:) within its attention weights
Layer 1, Posterior Responsibilities: The first layer uses the attention matrix w( to compute the posterior responsibilities and the relative positional encoding r'g to store the residual stream:
( (log πT" TT, (R)
J.
[ +50
T
2 2 k
otherwise
(R
)k.
10
Se
otherwise
S k:
2 m+1 m+2
2 m+1 m+2
T
1
0
0
(L) 81
R
(LYT
where log is applied element-wise, k = j-jj+I 1 is the relative position between token i and j shifted by The lookup table (²(")) biases attention to focus only on the first m relative positions (the lags).
The lookup table r2(") uses one-hot vectors e₁ to copy the computed attention weight for a specific lag into the corresponding dimension of the output vector, thereby storing the responsibilities for
'The attention sttentionbbased computation of responsibilities in the first layer is shared with the construction of D Angelo et al. (2025). However, our aggregation and prediction layers leverage RPR value embeddings to responsibilities in dedicated dimensions, yielding single-head construction.
5

---

Published as a conference paper at ICLR 2026
different lags in distinct positions. The attention score Eij for this layer is computed as per Equation Given that the input is one-hot encoded, i.e., h/ Cy the score becomes:
+01 if f m
TT (1) Eij = Ey (log πt" Cy; + Ey ij log x (vj Mi) +
(-ō1 otherwise
where we used that the RPE vector raj is constant with respect to the token values yi- For large 51, the softmax only attends to keys j such that their relative position
i - is within the
exples) range m). The causal attention weights A for j S i are given by: A Σj exp(e;yr)
x (v),)))) Ti
In the limit 61 - co, the behavior of the softmax changes based on the Ej-- #* (M)r-))) exp(e) "
position For the first token (i 1), the only valid key is itself (j' 1), which means the attention
is entirely self-contained, resulting in A = 1. For any subsequent token (i > 1), the terms in the denominator corresponding to lags & [1, m] are scaled by exp(51), while all other terms, including the diagonal (k = 0), are scaled by exp( (8)) and vanish. This yields the following limit for i > 1:
i-j E [m]
A
(MC-).8))
A
0
otherwise
Tr-1
Tr
where m min(i 1, m) For positions i > m, where the MTD model is defined, the attention mechanism thus computes A! = --(i- (j for the active lags (1 - S m). This expression is precisely the posterior responsibility (g) from Equation 6 where the lag is g = i i- j. For earlier steps (1 S m), the attention mechanism computes the natural counterpart by normalizing over the i 1 available lags. The resulting attention matrix (A)) has a specific banded lower-triangule structure. The output of this layer, h, is the attention weighttd sum over the concatenated value vectors. Given that the attention weights Ai compute the posterior responsibilities (g) (for i > m), and R embeds the lag = --1 as a one-hot vector eg, the output for a position >> m is:
i-1 Σ Au Concat(e,,
Σ (g) Concat(ew...: 9=1
(8)
This output is convex combination, weighted by the posterior responsibilities, of vectors that each concatenate two pieces of information: the one-hot encoding of the token at given lag (ey--g) and the one-hot encoding of eg of the lag itself (g):
Σxe H(g)em-p
Ege
Egg
71(1)
(m)
2
1
(1)
(2)
OH (m
Lag
Lag 2
Lag
In essence, the top part of (2 contains weighted sum of past tokens (which will not be used), while its bottom part explicitly stores the vector of posterior responsibilities with the value (g(g) for lag g stored the g-th position within this second block (i.e., (1) is first, followed by (2), etc.).
Layer 2, Summing Responsibilities: The second layer sums along the sequence the responsi- bility vectors computed in Layer for tokens at positions > This is achieved by setting the content-based attention to zero (W(22) 0) and the value-rpe matrix to zero (R(2) 0).
The mechanism relies on the content- content-position position interaction )T Crucially, the input vector h (1) Concat(ey, h ()) retains in the residual stream the one-hot embedding Cy. of the current token in its first q dimensions. The RPE table x²(2) is structured to interact only with this part of the vector, turning the dot product into fixed bias:
(2 A
Os, xd.
(R k.:
0 01xdi
if f otherwise
R
0 xm
1--2 11x9 O 01x(gtm)

---

Published as a conference paper at ICLR 2026
The attention score for the final query at i T therefore simplifies to:
(h% (-2z- 11 82 (e 11) -ō2 if ifl m
(2) Tj
0p+m
eTj
(hT (1)
Tod, =0
otherwise
In the limit 82 + oo, the softmax places uniform attention only on the keys where the score is not -O0. This results in uniform attention weights Ar; I/(T m) for Im + 1, T], and zero otherwise. The layer's output for the final token is therefore the exact average of the desired vectors:
T Ary Concat (h)) (T5)) m
h
Σ jamell
Co
70
T m
the m-dimensional sub-block at positions 2q++1 to 2m+m of hT contains the averaged responsibility vector T-m 2i=++1 Fy. which is exactly the quantity needed to implement the one-step MD estimator in Prop.
Defining F (7g (7)...)
Layer 3, Final Predictive Weights: The third and final layer uses the averaged responsibilities computed in Layer produce the final predictive weights, A, which correspond the one-step MD estimate from Prop. 2 This is accomplished by using the RPE table, R to perform a selective dot A
product. The query vector for the final token, le contains the vector of averaged responsibilities in the P sub-block of h.(e) The RPE vectors in (3) are constructed scaled one-hot vectors that align with this sub-block, effectively using the dot product to "read out" the corresponding averaged responsibility. Similarly to layer 2, to only have non-zero attention at the positions corresponding to the m lags, we use the one-hot embedding of the current token Ey: from the residual stream
h (2) Concat(eg, h (((2) add fixed bias 8g which drives the softmax to zero in the limit. The content- based attention and value-rpe are again disabled:
10)
(
on (3)
1:24
on
11:27:2m
W
Odzxda
(R 3
+83
Ck
if [m] otherwise
R V
Orxm
5g
4
0
m+
24
m
Here, eg is a one-hot vector in I that selects the coordinate corresponding to the k-th responsibility in the I sub-block of h and / is the learnable scaled learning rate. Visually, the RPE table R (3) is a sparse matrix of scaled one-hot vectors:
m
T
(0)
+0g
+03
1g
1
83
1
83
1
h1)
0
0
0
0
R
h 1:29
0
0
0
0
r
[Be]
Bea
Bem
0
0
0
hig m 2m
(2)
0
0
0
Since W( O xda the score simplifies to the content-position interaction (h )T and (3)
TT is constructed to be non-zero only for relative positions = T' - j + +1 E m]. For these lags, the RPE is a scaled one-hot vector that acts as selector, using the dot product to extract the corresponding averaged responsibility stored in l²)
ery (h
(3)
(h (0)) (8g1.) + (Be,)
+ōg +
Σr
Ei=m+ relg)
T-m
if E [m]] otherwise.
6z (e 1,)
53
The attention scores for the final token, eT,., becomes the scaled averaged responsibility:
er: = --0ōg.... - dz. tdg + Diamer (m)
+ + Dimmt1 7:(2)
+ + 3 Diam+r 7()
T
pos m-m+1
pos T
pos

---

Published as a conference paper at ICLR 2026
After applying the softmax and in the limit of large 63, the attention weights for the last token compute the MD estimate of the mixture weights Ag and place them at the correct positions:
B E v4
exp
lim Ar,T-g 63-000
lim Ar:: ...,... O,
R== exp Di-m+t (((
6g 0000
T-m+1
T-1
This final attention operation is the core of the estimation process. It computes the mixture weights X, which serve as the in-context estimates of token importance, thus realizing the mechanism of in-context learning this work seeks understand.
Output Layer: The final step of the construction is to apply the output matrix Wo to the final hidden state hT to produce the predictive distribution over the next token. The matrix Wo learns the known transition matrix π and selectively applies it to the embedding of the last token. The output of the third attention layer, fTT is a weighted sum of the hidden states from the second layer, where the weights are the estimated mixture weights Xg: h.)))
(where we drop the zero value-RPE block since R = Orxm). The first components of any Σド=
(3) h = - Xyhi
(R) due to the residual stream, simply contain the original input h2 Eng Conse- quently, the first q components of (3) are X-weighted combination of the one-hot embeddings of the relevant past tokens: (hT 1: = Ey= >g = Eg= Ågevr-g The full hidden state is h = Concat(h (0) hTp hr hT and the output matrix Wo E Raxda is structured to ignore all preceding blocks and operate only on the first components of the final block, h This is
hidden state h
achieved by storing the transition matrix, TT, in the corresponding sub-block:
(0)
(1) 0 xx (g
hT (2)
from h (3)
from RT
Wo-
O yxg
O (2g+2m)
[xg
Og (3g (3g-33m)))
Applying this matrix to the fully expanded final hidden state yields the predictive distribution:
F(y)T) | 0 h π T T(h
1:1
T
Σ
Xxeer
Àg (UT -g::)-
This final vector is exactly the predictive distribution from Proposition
completing the proof.
WHY ONE-STEP MIRROR DESCENT WORKS
theoretical analysis of the one-step MD estimator. We prove that single mirror descent update, initialized at the uniform prior, recovers a first-order approximation of the Bayesian posterior mean. This result provides formal justification for the effectiveness of the non-iterative estimator implemented by our construction.
We
One-Step MD a First-Order Bayesian Approximation: We establish theoretical connec- tion between the one-step Mirror Descent (MD) estimator and the Bayesian posterior mean. We show that their first-order Taylor expansions around the state of no evidence coincide up scalar constant. This result justifies interpreting the one-step MD estimator as principled approxima- tion the Bayes-optimal predictor, especially in low-data regimes. The analysis hinges on treating both estimators as functions of the log-likee likelihoodd gradient evaluated at the center of the simplex, = ())),, and expanding them around the point of no evidence, g = 0.
Theorem (First-Order Equivalence of the Estimators). Lei MM be the one-step MD esti- mator with learning rate d, and let Bayess(g (g) be the Bayesian posterior mean under the linearized likelihood. The two estimators are first-order equivalent at g 0 for
Learning-rate scaling via Lipschitz (smoothness) constant:
We established a first-order equiv- alence between the one-step MD estimator and the Bayesian posterior mean at g = (the "no- evidence" regime). For sequence of length T, however, the gradient norm |g] scales with 7 (see
8

---

Published as a conference paper at ICLR 2026
Lager
Layer
Lager
7 ddisent Ystandard
0.020
Constructed
7 Ticont.
KL. divergence
8.010
0.608
0.007
Traized (d disrentangeed
0.005
0.003
0.002
100
200
300
400
500
600
Sequence Length
Trained SeqLen 192 KL.00,0376
Trained (standar)
Figure 4: Comparison of Trained and Constructed Transformers. Left: Attention maps of the trained transformer (disentangled and standard) versus our theoretical construction (seq. length 64). Right (top): KL divergence the ground truth transition probabilities for the trained transformers, the constructed trans- former, and the one-step MD estimator across sequence lengths. Right (bottom): First-layer attention softmax Softmax(W: vs. true transitions matrix x' for a trained model.
App. raising the question of how to scale the learning rate of ÁD = softmax((1g) with T. Mirror Descent theory suggests choosing the learning rate inversely proportional the relative smoothness constant Luel (Bauschke et al., 2017). Because our MD update is derived from the negative entropy potential, smoothness is defined w.r.t. the KI divergence. Convergence requires 1) S 1/ Lel- We now bound this constant and find exactly the scaling implemented in the Transformer in Prop.
Theorem 2 (Relative Smoothness and scaling). At the center of the simplex À = (1/m..... /m). the loss f()) -((A) is Lmo-thhoth relative to the KL-divergence, with L 5 (T m) m². Consequently the stable step-size rule 7) S Lnd yields the asymptotic scaling 1j ((+) for fixed
Beyond the Local Approximation, The Implicit Regularization of Mirror Descent:
The first-order equivalence in Theorem [ holds only for short quences, where the log-likelihood gradient is small. For longer se- quences, neglected higher-order terms become significant, and the one-step estimator diverges from the Bayesian mean. Empirically, however, few additional Mirror Descent steps substantially reduce this gap (see Figure 5). In this regime, iterating MD to convergence yields the suboptimal MLE, while early stopping provides much closer match to the Bayesian mean. We propose that this effect arises from implicit regularization. Specifically, early-stopped MD approx- imately solves an entropy-regularized optimization problem of the form mİnA E(A) + H with the iterates tracking the correspond- ing regularization path (Suggala et al., 2018). Along this path, per- formance on par with the Bayes-optimal estimator is achieved for an appropriate choice of regularization 7 (see Figure B) Thus, early stopping is effectively equivalent to selecting this favorable point on the
MLE
10-1
Divergeen coerrenne 1Z
101
10a
Sequence Length
Figure 3: Regularized Es- timator. Comparison with Bayes and MLE estimators.
path, avoiding the suboptimal MLE.
6 EXPERIMENTS
We validate our main claims below; further details and additional experiments are in App. Band D Setup: We train 3-layer disentangled transformers Tasen with single head with learned relative

---

Published as a conference paper at ICLR 2026
positional and one-hot semantic embeddings well 3-layer standard transformers Tsundard (no disentanglement) with single head attention with standard parameterization given by query-Key- Value with learned relative positional and learned semantic embeddings. We sample a fixed q, m, and at each iteration we sample a set of { with B = 128 the batch size and X; ~ Dir(ax) and generate B sequences according to the MTD model. We train the model for 5 10 iterations using the Adam optimizer and MSE loss over the last token in the sequence for various sequence lengths.
Results for one-step MD: We plot the KL divergence with the ground truth transition probabilities between the trained transformer Taivem. and Taandard. compared to our theoretical construction Teomstr. the one-step MD estimator ÁMD and the optimal-Bayer estimator Bayes (we compute it via MCMC; see App. F.1 for various sequence lengths in Figure (right). We observe that both the disentangled and standard trained transformers match the performance of the theoretical construction and the one-step MD estimator: for small sequence lengths, the latter serves good proxy for the optimal Bayes estimator, validating our theoretical result in Theorem [ whereas for longer sequences becomes suboptimal. By inspecting the attention maps of the trained transformers in Figure 4 (left) we can see that they learn to extract the responsibilities r(g); as expected from our construction (for all three models, the locations of low and high attention entries, and the diagonal structure induced by the MTD order, are closely aligned). To further validate if the attention matrix in the first layer actually learns the ground truth transition matrix x² we plot the heatmap of the first attention softmax WATJ vs. n² as well as the average row-wise KL divergence between the two matrices Figure (bottom-right) more results in App. D For both the AD and Teumstr. we tune the learning rate / and 1y via grid search to minimize the KT. divergence (see Section B. for more details).
Results multi-step MD: To investigate whether deeper Transformers can learn to implement multiple steps Mirror Descent, we plot in Figure the the KL divergence for 5-layer trained transformer Tuain. compared to the multi-step MD esti- mator { MD. (i denoting the number of steps) across various sequence lengths, averaged over 3 seeds with error bars. We observe that the trained transformer closely tracks the perfor- mance of the 2-step MD estimator. We stress that this is a per- formance comparison, not convergence or optimality claim: we do not assert that the transformer converges to the 2-step MD solution. Rather, the experiment suggests that deeper transformers are capable of implementing estimators whose accuracy is at least comparable that of multi-step MD. No- tably, for longer sequences where the gap between the 2-step MD and the Bayes-optimal solution widens, the transformer performance exhibits some variance across seeds, occasion- ally falling below the 2-step curve. This hints that the learned estimator may not be tied to the 2-step MD and could, in principle, exploit additional structure. Extending our explicit construction the multi-step setting is nontrivial and we leave it to future work; however, the empirical finding that a 5 5-layer transformer suffices to match 2-step performance suggests that intermediate representations (e.g., the responsibilities computed in earlier layers) may be reused across steps.
7train.
Sayen
KL divergence
10-2
MD.
10:5-
500
1000
1500
2000
Sequence Length
Figure 5: Multi-Step MD vs. 5-Layer Transformer.
7 CONCLUSIONS
We set out to understand whether transformers can learn latent causal structures in-context, and what algorithm they implement to do so. To this end, we introduced a framework based on MTD that for- malizes -ikeonimpoetance estimation as in-context inference of latent variables. Our central finding is that transformers solve this task by implementing Mirror Descent. We gave an explicit three- layer construction that exactly implements one step of the algorithm and proved that the resulting estimator is a first-order approximation of the Bayes-optimal predictor, with learning-ral scaling 1 (1/TT)) matching the transformer's learned behavior. Empirically, transformers trained from scratch, both disentangled and standard, recover the constructed solution across predictive distri- butions and attention patterns, while deeper models achieve performance comparable to multi-step Mirror Descent. Taken together, our results extend the gradient- based interpretation of in-context learning from regression to sequential domains over discrete tokens, offering a new algorithmic lens for understanding latent-variable inference in attention- models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was partially funded by an unrestricted gift from Coefficient Giving, and the grant number 212111 from the Swiss National Science Foundation. Francesco D' D'Angelo is supported by the Google PhD Fellowships.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference Learning Representations, 2022.
Adrian Raftery André Berchtold. The Mixture Transition Distribution Model for High-Order Markov Chains and Non- Non-Gaussian Time Series. Statistical Science, 17(3):328 356, 2002. doi: 10. 10.1214/s/1042227993
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Thirty-seventh Conference on Neural Information Processing Systems, 2023a. URL http:///eeeeeeeeeeeet..oo
id= Id-1iMSqUuVg9
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2023b.
Heinz Н Bauschke, Jérôme Bolte, and Marc Teboulle. A descent lemma beyond lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research, 42(2):330-343 2017.
Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):667-77,, 2003.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transforeer: memory viewpoint. Advances in Neural Information Processing Systems. 2023a.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 20236.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 33:1877 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks of artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXiv:2303. 12712, 2023.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learn- ing in transformers. Advances neural information processing systems, 35:18878-18891 2022.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang- Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. In The Thirry-eighth Annual Confer- ence on Neural Information Processing Systems, 2024. URL
[ooummid 4 F
Arthur Conmy, Augustine N. Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL hetps:
//openreviens www.:nntt/oooommmmcmmmmm http://orrrmmm..89s991111777

---

Published as a conference paper at ICLR 2026
Francesco D' Angelo, Francesco Croce, and Nicolas Flammarion. Selective induction heads: How transformers select causal structures in context. In The Thirteenth International Confer- Learning Representations, 2025. URL https:/ httpss:/apeererieeet http:///ppeneeee.eettttetet net/forumidt
bnJgzAQjjff
Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. Advances in Neural Information Processing Systems, 37:64273-6433 64311,, 2024.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Haffield-Dodss, Danny Hernande7, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https://
Deqing Fu, Tian qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn to achieve second-order convergence rates for in-context linear regression. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. TtRphttss://ppeeett http///ppeeeeeeeeeeeeeeeeeeeeeee wwt. forum?
(d=L8h6cozcbr
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30590 2022.
Jiachen Hu, Qinghua Liu, and Chi Jin. On limitation of transformer for learning hmms. arXiv preprint arXiv:2406. 04089, 2024.
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self- attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint 235.13 2024.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 19689-19729. PMLR, 29-29 Jul 2023.
Yue M Lu, Mary Letey, Jacob A Zavatone-Vert Anindita Maiti, and Cengiz Pehlevan. Asymp- totic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, 122(29):e2002599122 22, 2025.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representutions, 2024.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.04161, 2024.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zeulemmyer. Rethinking the role of demonstrations: What makes in-context learning work? In Yoay Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 11048-11066 1164, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
Arkadij Semenović Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
12

---

Published as a conference paper at ICLR 2026
Timothy Nguyen.
Understanding transformers via n-gram statistics.
arXiv preprint
arXiv:2407.12044 2024.
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning, 2024. URL
net/ Forum?idie J H
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfiel-DDodde Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. In The Twelfth International Conference on Learning Representations, 2024. URL
net wtt. feteeummmmissss.... j91H120
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learming Research, 22(1 22(1):3463 04146633-44790
Adrian E Raftery. A model for high-order markov chains. Journal of the Royal Statistical Society Series B: Statistical Methodology, 47(3):528-539, 1985.
Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vard- han Makkuva. Transformers on markov data: Constant depth suffices. arXiv preprint arXiv:2407. 17686, 2024.
Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Michael Eli Sander, Raja Giryes, Taiji Suzuki, Mathicu Blondel, and Gabriel Peyré. How do trans- formers perform in-context autoregressive learning In Forty-first International Conference Machine Learning, 2024. URL URt.ttps::
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth. In International Conference on Machine Learning, 2024. URL https:
//api. www.semtsntniooaaa....eeccmcaaae wwwplaarrrrr//oo wwtp:oooooI.1::127555555.
Adam Shai, Paul M. Riechers, Lucas Teixeira, Alexander Gictelink Oldenziel, and Sarah Marzen. Transformers represent belief state geometry in their residual stream. In The Thirty- Anmual Conference on Neural Information Processing Systems, 2024. URL BRLL/tttss://ppeeee////
feet//orum?id=-IBBTe
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. In North American Chapter of the Association for Computational Linguistics, 2018. URL
www.ssetnnncccccccee... htct ar. wrtp/oorpssD::3322/22 37228115
Aaditya K Singh, Ted Moskovitz, Felix Hill, Stephanie C.Y. Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. In Forty-first International Conference on Machine Learning, 2024. URL https:
d=08rrx //openrevie eew.ewwwao www.:eet..orao : 71D5
Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar. Connecting optimization and regulariza-
tion paths. Advances in Newral Information Processing Systems, 31, 2018
Anej Svete and Ryan Cotterell. Transformers can represent n-gram language models. arXiv preprint arXiv:2404. 14994, 2024.
Aditya Varre, Gizem Yüce, and Nicolas Flammarion. Learning in-context SnS-grame with trans- formers: Sub-$n$-gramm are near-stationary points. In Forry-second International Conference on Machine Learning, 2025.
13

---

Published as a conference paper at ICLR 2026
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov. and Max Vladymyrov. Transformers lear in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174, PMLR, 2023a.
Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05855 058888 2023b.
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.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint arXiv:2311.0087/ 2023.
Oğuz Kaan Yüksel and Nicolas Flammarion. On the sample complexity next-token prediction. In The 28th International Conference on Artificial Intelligence and Statistics, 2025. URL https:
//operreveew. wtt.://ettttttommeooooaAAAAAAA
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeee functions? In International Con- ference on Learning Representations, 2020. URL http:://ppenneeviww.netttt wwwp:nrwwwwwwwweeeoo...eeeeeee
1d=ByxxMONtwrr
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett, Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.099: 09927, 2023.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization, 2024. URL.
ptttps://openrevieet www.
Nicolas Zucchet, Francesco D' Angelo, Andrew Kyle Lampinen, and Stephanie C.Y. Chan. The
emergence of sparse attention: impact data distribution and benefits of repetition. In The Thirty-ninth Annual Conference on Neural Information Processing Systemx, 2025. URL https:
//openrevies ew. wttpe/ewwwwww.weeeeeeeeeeeeeeeee formm? i ia Jh)MM44447
14

---

Published as a conference paper at ICLR 2026
A NOTATION
Throughout this paper, vectors (e.g., h, A), and uppercase bold letters for matrices (e.g., W, H, vector v is denoted by vs and the vector position in sequence of vectors H is written as h., The element in the i-th row and j-th column of a matrix A is Aiy. and its i-th row vector is A... We use and O denote vectors matrices of ones and zeros, respectively, with dimensions inferred from context. We use 6ş to denote a one-hot vector with a one at the k-th position; its dimensionality specified or clear from context. The set of integers (1...., m) is denoted by [m]. The probability simplex in Ih' is denoted by A The operator Concat(:.-) denotes the vertical concatenation vectors or matrices. For vectors
non-bold letters for scalars (e.g., 77, 0t), lowercase bold letters for πτ). The i-th element of
IRd and b E Rd,, their concatenation results in vector in &d.+ For matrices A E RdaxT and B Rdaxt xT with the same number of columns, Concat A, B) is the block matrix ( (dda+dn))) T Superscripts in parentheses, such as H used to index the layers of the Transformer. In the context of Transformer relative positional encodings, we use 1-based indexing for token positions i, E [T]. The relative position is mapped lookup table index k i j j+1 1. The approximations will be expressed using Landau notation, where a vector function fgg = (l]g signifies that [g) & Cllgg|" for some constant in a neighborhood of g 0.
A.1 DIMENSIONALITY OF THE TRANSFORMER CONSTRUCTION
The disentangled Transformer architecture results in hidden state that grows with each layer. The following table provides a summary of the dimensions at each stage of the construction. Note that the input to layer is (1-1), and its output (1) which is formed by concatenating the input with the result of the attention mechanism, (9) We use for the alphabet size and m for the MTD model order. The RPE value dimension is dæ m throughout the construction. Layer I uses the value RPE to store lag indices, while Layers 2 and 3 set R! O, so the concatenation appends m zero dimensions.
Table 1: Dimensionality of Hidden States in the Disentangled Transformer
Layer
Description
Attention Output (())
Concatenated Hidden State (1)
(Input)
Token Embeddings
h(0) E IR9
Responsibilities
h (1) Ra+m
(1) Concat(h (1) R2++m
2
Summation
(2)
2
h Concat (1) h(2)) R4++3m 3m
Weighting
h (3)) R4++3m
h Concat(h h(3) R89+6m 6m
Final Output Matrix
Wo E (8g+66))
B EXPERIMENTAL DETAILS
This section provides additional details on the experimental setup, including the hyperparameter settings for all models and estimators used in empirical validation.
B. HYPERPARAMETER TUNING
For the one-step Mirror Descent estimator (AD and our theoretical Transformer construction (Toosun) the learning rate parameters S) and 3 were not fixed but were tuned to optimize perfor- mance. For each sequence length evaluated, we performed a grid search over a range of potential values for and 3. The value that minimized the KL divergence the true Bayesian posterior mean was selected for the final comparison plots. This ensures that both methods were evaluated under their optimal conditions.
15

---

Published as a conference paper at ICLR 2026
B.2 PARAMETER SUMMARY
The following table summarizes the key parameters used in our experiments.
Table 2: Summary of Experimental Parameters
Component
Parameter
Value
Data Generation
MTD model order (m)
3,4,5
Vocabulary size (g)
Sequence Length (T)
Varied (64 to 1984)
Dirichlet Prior (at)
Uniform (0g for all g)
Transition Matrix (x)
Rows from Dirichlet (a 1)
Mirror Descent (MD) Estimator
Learning Rate (1))
Log grid: [10-5,10- 10-*1, 1000 points
Constructed Transformer (Tesastr.)
Large Constant (61, 62, 65)
[100.100.100]
Scaled Learning Rate (3)
Log grid: [10- 5, "1, 1000 points
points
Trained Transformer (Taisom..
Architccture
3-layer 5-layer Disentangled Transformer
Attention Heads
Concatenation
True
Semantic Embeddings
one-hot
Relative Positional Encodings
Learned
RPE value dimension
Emhedding Dimension
QK parametrization
False
Value matrix
False
Head output projection matrix
False
Optimizer
Adam
Learning Rute
x 10-3
Schedule
Batch Size
128
Training Iterations
x 10a
Loss Function
MSE the last token prediction
Trained Transformer (Tatandara.)
Architecture
3-layer Standard Transformer
Attention Hcads
Concatenation
False
Semantic Emheddings
Learned
Relative Positional Encodings
Learned
Embedding Dimension
QK parametrization
True
Value matrix
True
Head output projection matrix
Tru e
Optimizer
Adam
Learning Rate
x 10-3
IR Schedule
constant
Batch Size
128
Training Iterations
10a
Loss Function
MSE on the last token prediction
MCMC for Bayes Estimator
Sampler
Gibbs Sampling
Burn-in Iterations
[200]
Number Samples (K)
[2000]
16

---

Published as a conference paper at ICLR 2026
C RELATED WORKS
Induction heads and interpretability The emergence of ICI., well as the general ability of transformers to implement algorithms, has been linked to the formation of interpretable com- putational circuits [lhhae al. (202I) such induction heads (Olsson et al.. 2022). which are woven into sparse attention patterns Zucchhe et al. 2025). The development of these circuits is not monolithic; rather, they emerge in phases through the interaction of simpler subcircuits. Singh et al. (2024, for instance, use a causal framework to identify the key subcircuits whose interplay leads to the sudden formation of induction heads during training. This emergence itself also critically de- pends on the training data; specific distributional properties, such as burstiness and class imbalance, have been shown to be key drivers of this capability Chan et al. (2022): Zucchet et al. (2025). While much of this understanding comes from reverse-engineering circuits in pretrained models Commy eet al. (2023), a parallel line of research aims to create transformers that are interpretable by de- sign, for example by training models that can be directly decompiled into human-readable programs Friedman et al. (2023).
In-Context learning and gradient descent Following initial empirical observations of ICL in trans- formers (Brown et al. 2020), a significant line of research has sought to understand its underlying mechanisms. Harly work demonstrated that transformers can learn simple function classes like linear models in-context (Garg et al.) 2022). This led to the hypothesis that transformer layers effectively implement optimization algorithms, with several studies showing they can perform com- putations analogous to gradient descent for in-context linear regression (Akyürek et al. 2022, Bai al., 2023b Von Oswald al. 2023a b). This gradient- padeent-bbsed view has been extended to higher- order algorithms (Ahn et al. 2023; Fu et al., (2024) and given theoretical grounding, with proofs that gradient flow converges to a transformer that has learned the in-context task (Zhang et al. 2023). From statistical learning perspective, this process has been formalized as "algorithm learning" where generalization is guaranteed by the algorithmic stability of the Iearned procedure (Li et al,, 2023). From Bayesian standpoint, Zhang al. (2024) formalize ICL as Bayesian model averag- ing over latent concepts, providing generalization guarantees that complement the algorithmic view. Crucially, the emergence of this behavior is not guaranteed; it depends on sufficient pretraining task diversity (Raventos et al., (023). This algorithmic paradigm also extends to linear autoregressive processes, where transformers have been shown implement gradient descent step to learn the transition matrix in-context (Sander et al.] (022)). In a complementary direction, Lu et al. (2025) develop an asymptotic theory for in-context learning by linear attention, providing exact analytical characterizations in the large-dimensional limit.
In-Context learning, Markov chains and n-gram models Our work is closely related to the lit- erature analyzing ICL for sequential probabilistic models like n-grams and Markov chains. Foun- dational work by Yūksel& & Flammarion 2025) established formal generalization bounds for next- token prediction on Markovian data, analyzing its sample complexity. Regarding transformers, sev- eral mechanistic studies have investigated how they implement learning algorithms for these models. For bigrams, transformers have been shown to develop induction heads that function like associative memories (Bietti et al. 2023b) and accurately compute posterior probabilities from statistical cues (Edelman et al. 22224.. For first order Markov chains. Nichani et al. (2024) demonstrated that trans- formers learn the causal structure with gradient descent and implement induction heads to estimate the transition probabilities in-context, effectively implementing a Bayes-optimal estimator. This analysis was later extended to higher-order chains (Chen et al. 2024). while the work of D D'D'AAn A'gelo al. (2025) shows that transformers can even learn to select the correct Markov causal structure at inference time. Further theoretical results have explored the transformer loss landscape in this set- ting (Makkuva et al. 2024), characterized in-context n r-grams as star-stationary points (Varre al., 2025). and shown that constant-depth transformers are sufficient to learn k-th order Markov chains (Rajaraman al. 2024).
Transformers and sequential models Beyond specific learning algorithms, broader line of work has explored the fundamental capabilities and limitations of transformers as sequential models. On the foundational side, transformers have been shown to be universal approximators of sequence-to- sequence functions (Yun et al., 2020) and even Turing- complet under certain assumptions (Pérez et al. 2021), while recent work has investigated the trade-off between depth and parallel computa- tion Sanford et al.) 2024). In terms of representational power, transformers with sparse attention have been shown be capable of exactly representing any n-gram model (Svete & Cotterel] 2024).
17

---

Published as a conference paper at ICLR 2026
However, this expressive power has limits; for instance, transformers may be less effective at learn- ing certain Hidden Markov Models (HMMs) compared to RNNs (Hu et al. 2024). Interestingly, when investigating the internal representations that enable inference in HMMs, Ms, (Shai et al. (2024) showed that transformers maintain interpretable belief states that are linearly encoded in the residual stream. A strength of transformers is their ability perform in-context model selection. has been demonstrated that a single transformer can adaptively choose between different base algorithms or even qualitatively different tasks (e.g., regression vs. classification) based the prompt (Bai al. 2023a), effectively selecting between different function classes in-context (Yadlowsky et al.) 2023). While high-level n-gram statisties can approximate transformer predictions, the mechanism for how the correct "rule" is selected in-context remains an open question (Nguyen 2024).
D ADDITIONAL EXPERIMENTS
In this section we repeat the main-text experiments for different MTD orders, comparing the trained and constructed transformers: addition to the m = case in the main text, we report results for m = 3 and m 5. In Figure 6 we plot the KI. divergence to the ground-truth transition probabilities across sequence lengths for the trained transformers, the constructed transformer, and the one-step MD estimator (orders and 5). In Figure 7 we instead compare the learned first-layer attention (softmax) to the true transition matrices for orders 3 and 5, with the average row-wise KL divergence reported directly in each panel. Finally, Figure reports attention grids for the trained and constructed transformers (disentangled) at sequence length 64 for MTD orders m = 3 and m 5, analogous to the attention maps shown in Figure (left) in the main text.
train.
train.
0.020
Mo
0.010
divergence
divergence
0.008
0.008
0.005
0.008
0.004
0.003
0.096
0.002
200
4.00
600
800
1000
200
400
600
800
1000
Sequence Length
Sequence Length
ground-trute We report the KI. divergence to the ground truth transition probabilities for the trained transformers, the constructed transformer, and the one-step MD estimator across sequence lengths Left: order 3. Right: order 5.
Figure 6: KI. divergence
18

---

Published as a conference paper at ICLR 2026
Tricned Suglen KL-1
Trannd Saglon 254] KE
Thera
Thera
Loat
CR
1M0
cef
6M
CC
01
Trazze Seqlan 384| XZ=11200
Trazed Segion 512 X2-110004
Tree
Softia Apselte
Tree
Softce Anseie
Traced degles 6911 KL-11130
Trxzed Segten 744 Ki-d ceds
Tres
Tree
Softce Ansece
AMI
6K
427
K27
4:0
627
A20
CP
6c
Thared Sequin 328 RL-2
Thared Seaten 256 KD+0.0324
Sofmes inter
Triired Segten 384 R0-2.0084
Trxined Sesten 512| (2-6
CM
GH
C21
CP
CH
C2T
037
G2
C3
517
Thaired Sezion 640 KL-5.0353
Thained Secion TEB KLet.033e
214
AA2
011
341
Figure 7: Layer-1 Attention Softmax vs. True Transition matrices for orders 3 and 5. For both orders (m 3 top 3 rows, 5 bottom 3 rows), each panel reports the learned first-layer attention (softmax of W;AT) alongside the true transition matrix; the average row-wise KL divergence is reported in the panel title. Sequence lengths increase left-to-right, oop-oo-botoom.
19

---

Published as a conference paper at ICLR 2026
Layer
Layer
Layer
Constructed
Trained (disentangled)
Layer 1
Layer 2
Layer
Constructed
Trained (disentangled)
Figure 8: Comparison of Trained and Constructed Transformers (attention grids). Top: Attention maps of the trained transformer (disentangled) versus theoretical construction (seq. length 64, MTD order
3). Bottom: Same top but for MTD order 5.
20

---

Published as a conference paper at ICLR 2026
E DERIVATION OF THE BAYES- OPPSSAL MTD PREDICTOR
Here, we provide the full derivation for the Bayes-optimal predictive distribution stated in Proposi- tion !
Our goal is to derive the predictive distribution for the next state, p p(Y++ vi, a),, given observed data prefix y (y)....,,.) We). The unknown mixture weights (A1,... Am) are assumed to be drawn from a Dirichlet prior distribution:
p( = Dirichlet(Ala)) = r Ogy T
II
[(ag)
9=1
The likelihood the observed data given the parameters is defined by the MTD model:
plui X)
II
p(ue vi
X)
II
Σ
An ho un)
k==+1
k=m+1
A==
Combining the likelihood and prior via Bayes' theorem yields the posterior distribution over the mixture weights:
p yi a) & plví x) a).
The Bayesian predictive distribution is formulated by marginalizing the single-step prediction p(Y++1 = vi X) over this posterior distribution of A:
/
p( ji jiainilliiie a))
p(Y++1 j A) p(A 3 a) dA, AA
where Am-1 is the probability simplex. The single-step prediction is simply the MTD model defi- nition:
m PY = ==liiiiiii Σ
Ag (y++1-9:))-
9=1
Substituting this into the integral gives:
m
p( = vi,
Σ
Ao (V++1-91))) pla ví. dx.
Am-1
By the linearity of expectation (and integration), the integral and the finite can be interchanged:
m
(S >g p(A yl. a) az) dd
p(Y+1 = 1
Σ T (1++1-9:3)
We recognize the term in the parentheses the definition of the posterior mean of the parameter Àg:
f Amm
Bayes = E[>y
vi ad)
p(A ví. a) dA.
This substitution yields the final form of the Bayes-optimal predictor, completing the proof:
p(Y+1 = 1o.ilvi.o vi
Σ iBaes T(1+11-9:5).
E. THE STRUCTURE OF THE BAYES-OPTIMAL ESTIMATOR
While the posterior mean is intractable to compute, its structure can be derived exactly. In particular it can be shown that the MTD posterior is finite mixture of Dirichlet distributions, and from this, can derive that the mean of the posterior preserves the classic 'add-constant' structure of conjugate Bayesian models, where the unobserved data counts replaced by their posterior expectation.
21

---

Published as a conference paper at ICLR 2026
Proposition 4 (Posterior Mixture of Dirichlets). Given the MTD likelihood L(X) plví X) and a Dirichlet prior p(> x) = Dir(A a), the posterior distribution is a finite mixture of Dirichlet distributions:
p(A ai =
Σ =E(1,..,mK'-
- a a++k(z)))
(9)
where 2 (2m+1,.... 21) is a latent assignment path, k(z) is the vector of counts of each lag in path and (T) p(Z z yi α) are the true posterior probabilities of the latent paths.
Proof. We begin with Bayes' theorem for the posterior distribution:
p a) x plyl X) p(> (x).
The central idea is to express the observed-data likelihood, p(yf | X), by marginalizing over all possible latent assignment paths z. A path x (Am+1..... 2)) specifies which lag was used at each step k.
plvi 1A)
Σ
plyi A).
=E(1...,m)'--
The joint probability of the data and specific path 2, known the complete- data likelihood, is given by:
plyl. 2 A) = II
plun 2k y wi-1. A)
k=m+1
II p(=x X) plut I yo 1x, k=m+1
=x.
II Azy - k=m+1
We can group the terms that depend on and those that do not. Let kg(2)) El Ex=m+1 I(2x be the number of times lag is used in path z. Then:
plvi 1
II
(=)
II
π T3) 2x 8k)
k=m+1
Let Ply 2) = -=m+1 11 ग. (ilk- yk), which constant with respect to X. The prior is given by m Ag - where B(xx) is the multivariate beta function. Substituting these
p(A a:) = B(a))
into the expression for the posterior:
V II
(+) )
m
p(A a) &
EP(NE
B(a) II -
B(a)
g=l
1=1
m
x Σ
Ply 2)
(x)-+)
1=1
We recognize that the term II, A (b)))( is the kernel of Dirichlet distribution, Dir(A a + k(2)). We write it ((x k(z)) Dir(A a + k(2)). Thus, the posterior takes the form of weighted sum Dirichlet densities:
p(A yi, a)
Σ (()) Dir a + k(=))
where the mixture weights T2)) are the normalized coefficients, which are precisely the true poste- rior probabilities of the latent paths, p(z = z yf, xx.. This completes the proof.
From this mixture structure, derive exact identity for the posterior mean.
22

---

Published as a conference paper at ICLR 2026
Proposition 5 (Bayes Mean Bayesian posterior mean, Bayes E[A yi ox), are given by:
Coosstnttto-Expeetee--e
The components of the
2g Ez K 2)
Bayes
(10)
Go m) Ezag. 0g. and Ez (2)) is the posterior expected number of times lag was used, which is computationally intractable.
where ao = =E
Proof. The posterior mean is defined by the integral of À over its posterior distribution:
Bayes
p( ví. a dA.
Am-1
We substitute the mixture-of- form of the posterior from Proposition
Bayes
[x(e)
x2))
Dir(A a + k(z))
dA.
Amm-1
Since the summation is over finite set of paths 2, can swap the integral and the summation:
X Bayes
=x())
A Dirrir(A
k(=))
The term inside the parentheses is definition of the mean a Dirichlet distribution with parame- ter vector [3 == a + k(2). The mean of a Dir(B) distribution is the vector B/Bu, where Bu EE Bj- In our case, the sum of the parameters is:
Σ + k=)=)) =
Σ ogg +
Σ Kg((z) = og + (m-
Therefore, the inner integral evaluates to the vector
a+*((x) (t-(t-m) Substituting this back:
+ k(2)
X
Bayes
Ex(a)
Og + (t m)
This expression is an expectation over the posterior distribution of latent paths, Z T(2). We can write it as:
Ez~T a + k(Z)
Bayes
au (t mm)
By the lincarity of expectation, we can take the expectation inside for each component g:
E2~= (ag + (2) ag + (k,
Bayes
9
(1 m) Okg (m-m)
This reveals the "udd--oottantt---ppccdd-oooo structure the estimator, completing the proof.
F APPROXIMATIONS FOR THE MEAN OF THE POSTERIOR DISTRIBUTION
In this section, we outline the methods used approximate the mean of the posterior distribution over the mixture weights A.
F.1 MARKOV CHAIN MONTE CARLO (MCMC)
Since an analytical solution unavailable, we approximate the Bayesian predictive distribution using samples from the posterior distribution generated by Markov Chain Monte Carlo (MCMC) method, specifically Gibbs sampling. Gibbs sampling is well-suited for this problem because, while the full posterior is intractable, the conditional posteriors of the parameters and latent variables simple to sample from.
The procedure involves augmenting the model with latent variables Z (Zm+1..... Zi). where Z = g indicates that lag was used to generate the transition to Yg. The Gibbs sampler iteratively draws from the two full conditional distributions:
23

---

Published as a conference paper at ICLR 2026
1. Sample Latent Variables Z given Parameters A: For given parameter vector A (R-1) from the previous iteration, we sample each latent variable Z₂ for s E {m + ...... f from its categorical conditional posterior:
1 (k- 1) n (yss 1al--911 Us)
P = yi, X 1
(11)
En= X/2 (Vo- Us))
This provides a complete sampled sequence of lags, (1)) (zum11...... 20
(8)
2. Sample Parameters À given Latent Variables Z: Given the sampled lags 2(²), the Dirichlet prior is now conjugate to the complete-date likelihood. We first count the oc- currences of each lag, "'g = Ei-m+1 T(z! g). The full conditional posterior for À is Dirichlet distribution, from which we draw the next sample (k):
p(> vi (Ax), a) Dirichlet(A 01 m......, Om mmm)
(12)
After running the sampler for Kisal iterations and discarding an initial burn-in period, we obtain set of K samples, {A(1) (2) {k)), that are approximately drawn from the true posterior
p( vi, a) shown in Equation
The integral in Equation 4 is then approximated via Monte Carlo average:
p( 11+1=1lvi = K R
Σ
p( = jivl. (k))
1
Σ
ग. (1+++
(13)
This estimate converges to the true Bayes optimal predictive distribution as K co. It represents the theoretical performance limit for inference under the MTD model assumptions, providing a gold-standard benchmark against which other estimators can be compared.
G ALGORIHTM FOR MAXIMUM LIKELIHOOD ESTIMATION OF MTD
Maximum Likelihood Estimation (MLE) for the Mixture Transition Distribution (MTD) does not admit closed form solution, therefore iterative optimization algorithms are required. This section details two iterative algorithms suited for this task. We first review the Expecation-Maxxiization (EM) algorithm, standard and widely-used method for latent variable models. We then present Mirror Descent (MD), an alternative optimization framework that is central to our work.
G.1 EXPECTATION- MAXIIIONTNAAIIMAAIIINO (EM) ALGORITHM
The popectaaionnMaximizattee (EM) algorithm is widely-used iterative method for finding Max- imum Likelihood estimates statistical models with latent variables. In the context of the MTD model, the latent variables correspond the specific mixture component responsible for generating each observation. The algorithm alternates between two steps: the Expectation (E) step, where it computes the expected likeliioooood with respect the posterior distribution of the latent variables. and the Maximization (M) step, where it updates the model parameters to maximize this expected value.
Given the latent variables Z = (Zm+1..... ZT), where each 2, E {l,..., m). The variable 2, indicates that the 9t mixture component (corresponding to lag g) was responsible for generating the transition to Y, at time t. The complete data are (3. 2).
The likelihood of the complete data (y TT 1). conditional on yी" is given by:
T
P Ply TT
II
P(ye 2t 1 -1A)
t==+1
T
II t=
P(Zz = 24 vv A)P(MI 2, 21, 21,11- -1;
24

---

Published as a conference paper at ICLR 2026
Under the MTD model assumptions:
P(z = : Xg
P(ye Z: - 9, vi ; >) = π (V)- 3/1))
Thus, P(y. Ze 9 yi ; ;) > wt). The complete data likelihood becomes:
P(y 2m+1 "y" A)
II (==++1
Ao i. (1t Wt).
The complete data log-likelihood, (o 2) = log P( 2 "i" A), is:
(o():
Σ
log(>2, Te. (1))
t=m+1
11 = 9) V:))
(14)
(= g=1
where I(.) is the indicator funetion.
G.2 EM ALGORITHM STEPS
Let A (R) be the estimate of À at iteration
E-Step (Expectation) The E-step computes the expectation of the complete-data log-likelihood equation 14 with respect to the conditional distribution of the latent variables Z given the observed data y and the current parameter estimate (A),, This expectation defines the Q function:
Q(A A (*)) = Ez (44) ll (A: y. 2)]
E
Σ
I(Z 9) log (V0 90)
y. (1k) (b)
[=m+1ge
T
B[T 9) y. x (k) log (V)-
1= g=1
The core computation is the posterior probability (responsibilitty of Ze = g:
2t E| g) y, )*) P(z = g y, (k)).
Due to the MTD model structure, future observations yf+1 are conditionally independent of Ze given y.. Thus, the posterior probability simplifies:
IP g y. (k)) P(Z: = |
Using Bayes' theorem:
(g) P(Z. =g gl (
P(y Ze 9. vi -i
= 9 y ((()))
Ply V (()))
x (v: 3/ (A)
M= P(ys, Z, h vi (k))
X T (Vt--::11 Vt
(15)
(k) N ( 31
The E-step involves calculating these responsibilities %t (8) (g) for all E fm + 1,..., (1. The Q function is then:
......T T) and g E
Q(A (())
Σ
(k) Yt
(g)g Ag + log (og
(16)
1=1
25

---

Published as a conference paper at ICLR 2026
The M-step finds the parameter values X that maximize the Q function equation 16 subject to the constraints 2g 2 0 and Σg=1 Ag = 1. This gives the updated estimate
M-Step (Maximization)
A (k+1)
( = arg max Q(A | A (1)).
Since the terms log (8)- Vt) do not depend on X, we maximize:
m
m
T
f(A)
Σ
(g) log Ag - Σ
Σ
't'
(x) (g)
log
t= g=1
9=1
M=m++
Let Cg = -E Tt
(g). We maximize f(A) = -Ey=1
Cg log Ag subject to Eg= Xg = 1. We
use a Lagrange multiplier
Σ Ag -1
((), = Σ
Co log Ag - 1
1=1
9=1
Setting partial derivatives to zero:
dL
Cg
Cg
/ 1=
Ag
2gg
Xg
m
m
Ag
Σ Ag 1.
g=1
g=1
Substituting X Cg/P p into the constraint yields μ -hol Ch. Therefore:
21=m+1 1tt (A) (g)
Co
Ag Eis Ch Em= Ex=m+1 9t' (h)
The denominator simplifies as Ex= Et=m+1 Ye
(1)(1) = Ex-m+ 2h=1 7t (h)
(1)
(1)
W 1 T- m. The M-step update rule is thus:
E Et=m+1 7t (g)
(1)
X 1
(17)
T-m
G.3 SUMMARY OF THE EM ALGORITHM
The FM algorithm for estimating the mixture weights À in the MTD model, assuming known tran- sition matrices Tg. proceeds as follows:
1. Initialization: Choose initial weights ( = (2 (0)
( such that A 2 0 for all
g E (1....,m) ,m] and Σ ( 1. Set the iteration counter k = 0.
2. E-Step (Expectation): Compute the responsibilities 7:
(M) (g) for each time point e {m + 1,..., .T) and each mixture component g E (1,-.., m. using the current parameter estimates (k))
X (y1
(g)
(18)
Ch An (Vr Vo)
3. M-Step (Maximization): Update the mixture weights X (k+1) {I...., m. using the computed responsibilities:
for each component
Σl 7t (1)(g)
(19)
T m
4. Convergence Check:
If the change in the parameter estimates (e.g., [|{(+)) - X (D
or the change in the observed data log-likelihood (e.g., ((x(+1)); {(> (k): where {(>: y) is the observed data log-likelihood) below predefined tolerance c, stop the algorithm and return Å = A (R+1) as the estimated mixture weights. Otherwise, set k t- k + and repeat from Step 2.
26

---

5): Published as a conference paper at ICLR 2026
G.4
MIRROR DESCENT
This appendix provides detailed derivations for Equation ( and Proposition
G.5 DERIVATION OF THE EXPONENTIATED GRADIENT ALGORITHM
The Exponentiated Gradient (EG) algorithm is a specific instance of the Mirror Descent (MD) online optimization algorithm. We apply to the problem of maximizing the MTD log-likelihood function, {()), with respect to the mixture weights À (A1...., Am). The weights are constrained to the probability simplex, Am- {A m Eg -, 0}.
The general MD update step iteration & linearizes the objective function around the current esti- mate (k) and adds a regularization term. The next iterate, ( is found by solving:
A (k+1) arg max { (VE(A (4)) >) - 1) AD AEAm-1
(20)
where 17 > O the learning rate, (E() is the gradient of the log-likelihood evaluated at A ( and De is the Bregman divergence associated with potential function V. For optimization the simplex, the standard choice is the negative entropy potential, (((A) = mm Xg log Xg- The resulting Bregman divergence is the unnormalized Kullback-Leible (KI.) divergence:
A
D (A
g log
(21)
(R)
Ag
Therefore the optimization problem in Equation 20 becomes:
m
A (++1) arg max AEA,
(V((A A)
A log
(22)
(E)
1
Xg
To solve the optimization problem in Eq. 20 we form the Lagrangian with a multiplier μ for the constraint Eg Ag 1:
((EE)) ),
Σ
Ag log
Ag
Σ >gg-11 {
(23)
Ag
Setting the derivative aL/0>g to zero yields:
>g (k)
17 ()
log Ag
+
p 1=
log
2
1 (x
-
(k)
Xg
Ag A exp(1 (x(((()) exeeppppp 1).
The term exp(-ı 1) serves a normalization constant to ensure Eg Ag 1. This leads directly to the EG update rule presented in Equation
(k) exp(n (x((A(1)
A (k+1)
repeated)
exp(1 VA((A()))))
G.6 DERIVATION OF THE ONE-STER MTD ESTIMATOR
We now prove Proposition by specializing the EG update the MTD model and evaluating the uniform prior (0) (1/m..... 1/m).
27

---

Published as a conference paper at ICLR 2026
Step 1: MTD Log-Likelihood and its Gradient. log- ligliikelhhod is:
Given an observed sequence prefix yi. the MTD
((A) = logp(yi X)
log
(24)
k=m-1 1
h=1
The g-th component of gradient
Vx((A)g = (E(A)
π (yk- g: ykk
(25)
8bg
h= A (3 - ykk)
k=m-1
Step 2: Evaluate Gradient at the Uniform Prior. We evaluate this gradient at A
(1/m,.... 1/m):
T (ikk g. 1kk))
k=m+1 h==1 (1/m) m (yk- yk) V)E(A(0))
(yk-ga go yk
(26)
Σ π (1k- 2k)
Amm++11
We recognize the term inside the summation as the posterior responsibility of lag g under the uniform model:
(x( = p(Zz = 3i (0))
(27)
p(uk Yk- )p( Zs = 90) EiPP Uh b) Zx - 111101)
(28)
T Tyk- -g Uk) (1/m)
(29)
Σn (ykkhh h. 4k) (1/m)
T(Kk-9 g1 21h
(30)
Σ πT (8k h ykk
Thus, the gradient at the uniform prior scaled sum of these responsibilities:
Vxl(A m
kk
umil
(g).
(31)
k=m+1
Step Apply the EG Update Rule. into the EG update rule (Eq. 5) to find !
Finally, we substitute {!
1/m and the derived gradient
X exp (1) vx((A (n)
(0)
A MD 9
x
ΣØ 110 exp (17 V)e(A(01),)
(1/m) exp 17 m Ek=mt1 unif (g)
(m exp () m Ek-m+t 7it unif (j)
m
exp
1pm Ek-m+1 (m)
(32)
E exp
1/ El Ex=m+1 7kk (j)
This completes the proof of Proposition 2
H ONE-STEP MD AS A FIRST-ORDER BAYESIAN APPROXIMATION
In this section, we formally establish theoretical connection between the one-step Mirror Descent (MD) estimator and the true Bayesian posterior mean. We demonstrate that their respective first- order Taylor expansions around state of "no evidence" identical up sealar constant. This
28

---

Published as a conference paper at ICLR 2026
result provides rigorous basis for understanding the one-step MD estimator principled approx- imation to the Bayes-optimal predictor, particularly in low-data or low-signal regime.
The analysis hinges treating both estimators as functions of the log- logelikeldhood gradient evaluated at the center of the simplex, g V Vx ((x()) and expanding them around the point of no evidence, g = 0. The approximations will be expressed using Landau notation, where vector function (g) ([]gg)))) signifies that [ I ll | for some constant C' in neighborhood of O.
Proposition 6 (First-Order Approximation of the One-Step MD Estimator). Let (4D be the one- step MD estimator defined as (g) softmax(ng) viewed as function of the log oogliikelhood gradient g. Its first-orerr Taylor expansion around g given by:
1 m Okk 9) (11991)),
(33)
where 9 = - -==19 m ( = (1/m...., 1/m) and Proja is the operator that projects vector onto the hyperplane of vectors that sum zero.
By. In vector form, this is (g) 1 2 Proja(g) + O(llg)?), where
m
Proof. The one-step MD update is given by the softmax function, XM > -ب- exp(ngs) Since the exponential function is analytic, the softmax function is also analytic in g. and its Taylor series expansion around g = 0 exists. The expansion given by:
exp(ugs)
MD x MD (0) + J/m (0gg + O(1gg)
where A/DD (0) is the Jacobian matrix of \ evaluated g 0.
Zeroth-Order Term: At g = O, the estimator evaluates to the uniform distribution, which the prior mean A (0).
exp(O) Σ プ (X)
MD (0) x
m -
First-Order Term (Jacobian): The entries of the Jacobian matrix, Jay(g) = à X (g), are given
0g, by n AMD (g) (g)) Evaluating at g O, where (MD( 1/m for all j:
aAMD 0g, g=0
1
5kj 17 m
m
Assembling the Expansion: The k-th component of the expansion is (M = (MP(0) +
ビj=1 0g, 9j + O(1991?):
m aiMD
8kj m m
MD (g) =
11
9j + O(lgg1
==1
1
17
m
8k19
Σ
m
9j
+ +(lgg])
m
m
m
j=1
7. (gk g) + 0(1991²)
This completes the proof.
Next, we derive the corresponding approximation for the Bayesian posterior mean. We linearize the log-likelihood around the prior mean, (0), which allows for an analytical treatment of the posterior.
Proposition 7 (First-Order Approximation of the Bayesian Posterior Mean). Consider a Bayesian model with uniform Dirichlet(1 1) prior over X Am-1 and likelihoood linearized around
29

---

Published as a conference paper at ICLR 2026
the prior mean, {(X)
(() + (g. x The resulting posterior mean, X R has
first-order Taylor expansion around g = given by:
Bayes
g
+ m(m + 1) (gs -g) +(1gg]))
(34)
m
In vector form, this is Bayes (g) variance matrix of the prior distribution.
( Cov (A)g + O(1gg])), where Cov (A) is the co-
Proof. Under the linearized likelihood, the posterior density is p()g) & p(A) exp(((A)) &
expl(g (g. A)), where terms constant in À are absorbed into the normalization constant. The poste- rior mean is:
Å Bayes (g) JAm_ exp( 2) dn(A)
JAm_, exp( (g. A)) dj(A)
where (p(A) is the uniform probability measure the simplex Am-1 Let (g) be the numer- ator's k-th component and (gg) be the denominator. The function Å Bayes (g) is analytie, allowing Taylor expansion.
Zeroth-Order Term: At g 0, the exponential term is The posterior equals the prior, so the
posterior mean is the prior mean:
Sam_, Adμ(A)
Bayes (0) Sam_ 1dp(A)
-Bx-Di(1) = (
First-Order Term (Jacobian): The Jacobian entries O4*
are
( N₂ 3g,
0
(1)
Using the quotient rule:
0g.
1
ON
az
Z
NR
0gj
72
0gj
0gj
We find the required derivatives (g) and (g) differentiating under the integral sign, which is applicable here as the integrands are continuous on the compact domain Am-1:
az
0
du(a) =/a >g exp (Eard) )d(aa)
exp
0g;
0gj
Am
0Nr
0
Egids dp(l)) Am-
Axds exp [oirr) du(a)
d exp
0gi
0g;
Am-1
Now, we evaluate these components at g 0:
Z(0) = f ddp()) 1.
.N((0) f >a du() = E[Aa] = 1/m.
= -Sam A A)exp((0, 2)) du(()) 1 du(2)== = E[A,] = 1/m.
89: g=0
aNL Øgg
SAm_, expe(00 ())) = A du(A) E[Aedg].
g
Substituting these evaluated terms into the quotient rule expression at O:
1 E[A,A, (1/m) (1/m) E[A
E A B[A,] Cov so (y
0gj 0==0
12
This establishes that the Jacobian of the posterior mean is the prior covariance matrix.
30

---

Published as a conference paper at ICLR 2026
Assembling the Expansion: For Dirichlet(1) distribution, the covariance matrix is Cov Aạ) mox;=1
The k-th component the expansion is:
m²(++])
Bayes k
1
m (m 1)35 95 + O(llel?)
mokj
m
mgk Σ 9j j=1
O(lgg1))
m
m² (m +
1
m
(9x - g) O(1gg]²)
m² m²(m (m +
1
m m(m + 1) (gk 9) + (1gg)))
This completes the proof of the proposition.
By comparing the results of Proposition and Proposition 7 we arrive at our main result.
Theorem 3 (First-Order Equivalence of the Estimators (restated)). Let MD (g: be the one-step MD estimator with learning rate n, and let Å Ruyyessg (g) be the Bayesian posterior mean under the linearized likelihood. There exists unique learning rate 1| m for which the two estimators are first-order equivalent at =
Proof. Two estimators are first-order equivalent at g = if their values and their Jacobian matrices
are identical at that point. As shown in Propositions and 7 the first condition, X MD XMD(0; n) =
Bayes (0), holds for any 1). The equivalence thus depends matching their Jacobians.
From the proof of Proposition 6 the Jacobian of the MD estimator at g is:
JAD (0)
17
11
m m
From the proof of Proposition 7 the Jacobian of the Bayesian posterior mean is:
(1-111")
Jans (0) m(m +
Equating the two Jacobians, J AmD (0) J Snye (0), requires their scalar coefficients to be equal, since the matrix factor is non-zero for 1:
1
m(m + 1)
Solving for yields the unique solution 17 ++
This theorem provides theoretical justification for the performance of the one-step MD estimator. demonstrates that this simple, non-iterative update is not merely a heuristic but a principled, first- order approximation of the Bayesian posterior mean under linearized likelihood.
Remark (The Role of and the Small-Gradient Assumption). The first-order equivalence es- tablished in Theorem holdds in the regime where ||g]] - 0, as this where the higher-order terms, Olllg))), are negligible. For many models, the gradient's 's magnitude, [gg], scales with the amount of data (e.g., sequence length 7), which appears to invalidate the approximation when the data size is large.
However, the one-step MD estimator, AD softmax(ng)) can remain well-behaved estimator even for large ||g| if is scaled appropriately. The term 77 determines the softmax behavior. If we set the learning rate be inversely proportional the signal strength, for instance r7T") O(1/T), the norm of the argument, lg]]. can remain bounded. This scaling prevents the softmax output from saturating and allows the estimator remain sensitive the information in g-
31

---

Published as a conference paper at ICLR 2026
On the Bayesian side, under the linearized likelihood, the posterior mean Rayye lacks scaling parameter analogous to 7). For large ||g]|, the posterior density exp( (g. 2)) concentrates sharply at the vertex of the simplex maximizing the inner product with g. In this case, the first-order approxi- mation from Proposition breaks down as the neglected O [ term becomes dominant.
The equivalence in Theorem should therefore be interpreted shows that for low-signal scenarios, the MD update is a principled approximation to the Bayesian one, and it provides a theoretically grounded value for 7) in that regime. The empirical success of the one-step estimator for large 7 suggests that while the functional forms of the two estimators diverge beyond the first order, a properly tuned MD estimator may still serve as an effective proxy for the Bayesian posterior mean, motivating analytical frameworks beyond local Taylor expansions.
local consistency result at g 0.
I LEARNING-RATE SCALING VIA A LIPSCHITZ (SMOOTHNESS) CONSTANT
In our main analysis, we established a first-order equivalence between the one-step MD estimator and the Bayesian posterior mean in the regime of "no evidence," where the log-likelihood gradient = O. However, for a sequence of length T, the magnitude of the gradient, |g typically scales with T. This raises the question of how the learning rate of the one-step MD estimator, MD softmax((gg) should scaled with 7 maintain good performance.
In this section, we provide theoretical justification for scaling the learning rate as 1 ( We demonstrate that the negative log- eigllikelhhood function, viewed as loss function over the simplex, has gradient that is Lipschitz continuous with constant L that grows linearly with the sequence length T. Standard optimization theory suggests setting the learning rate inversely proportional to this Lipschitz constant, i.e., x L,L, to ensure stable updates. We use the same notation as the before:
Crg T V/);
t= m= ,,
g L ..... m.
and m the number of lags.
Assumption, We assume there exists a constant
Cmin >
such that for every / E {m + 1, .,..., T} and every E m.)
Ct.g = (B)- g. 91) 2 Cmin-
(35)
Because each Ct.g is conditional probability, we also have the trivial upper bound
Ctgg S for all g.
(36)
Under these assumptions the denominators that appear in derivatives are uniformly bounded away from zero, and global, uniform bounds on the Hessian are valid.
Lemma (Hessian decomposition). For {(A) Et log Σn )gct,a) the Hessian satisfies
C.C.
V'E(A)= -- E
c= Cm,m
CE
)gtr.g)) )²
tmm+1
Hence the Hessian of the negative log-likelihood S(2) = ((A) is
T C.C.T
V²((A)
Σ
)gc1,))
(=++1
Proof. Direct differentiation of É(A) yields the displayed formulas.
Write se(A): ;= Σn AgCt.g and consequently the Hessian of the loss
T
V²((A) = Σ
C.O.
(37)
8: (A)
(==+1
Each summand in equation 37 is rank-one positive semi-definite matrix.
32

---

Published as a conference paper at ICLR 2026
1.1 UNIFORM OPERATOR-NORI BOUND ON THE HESSIAN
We now derive a uniform bound on the spectral (operator) norm of V (A) that depends linearly on T
Lemma 2 (Operator-norm bound). Under the standing assumption equation 35 (and equation 36 for every X in the full simplex A - (> 2 0, Eg Xg = 13 we have
[())ll I (T - -m)
m
(38)
min
Proof. From equation 37 and subadditivity of the operator norm,
C.C.
|c.cc Tlop
V f())l =
8e (A)²
St {²(2
1==+1
1=++1
For rank-one matrix uu the operator norm equals ||
Hence
m
c..c7 llap
||call2
Zg
9=1
Using equation 36 we get the upper bound ||c, 1/ < m 1 = m for every t.
For the denominator, by equation 35 and since
g=1 2g
m
m
(A
2 2
Ag
Cmin Cmin
9=1
1=
Therefore for every t,
|c,c llop (A)2
m min
Summing over t m + 1,.... . yields
T
m
m
5
(T - m)
1=
min
min
which is the bound equation 38
1.2 IMPROVED BOUND AT THE CENTER OF THE SIMPLEX
Here show that evaluating exactly at the uniform vector X yields a strictly better constant: the spectral norm of the Hessian at A is bounded by (T-m) m². This gives less pessimistic Lipschitz constant and hence looser restriction on the conservative step-size n7.
Proposition 8 (Operator-norm bound the uniform vector). Let x (1/m,..., 1/m). For the loss f(A) = -(()) we have
|
(T (---mmm² m²
(39)
In particular, the gradient V Lipschitz X with constant L S (T m)m? and the conser- vative step-size choice 1) I 1/L" yields 1) = ((1/T) for fixed m.
Proof. Recall the Hessian decomposition (Eq. equation 37
C.C.
m
v()())
St (A)2*
C: =
s((A) = Σ
Ao Cr
(=m+1
==1
Evaluate at the uniform vector A Then
1
S
(x*))
Crg
Σ Ca.9
m
m
9=1
33

---

Published as a conference paper at ICLR 2026
Hence the t-th summand becomes
C, Sy (X+)
G
m²
C.C.
(St/m)2
S
Taking operator norms and using subadditivity.
T
m²
c.c. llap
St
1=
For each f, ||c,c,] llop = ||c, 11.. Observe the elementary inequality
m
2 = S? 1t.9
cc 1
9=1
Σ, at + Eisr ajaj 2 E, at for nonnegative a, Using this which holds because a inequality we obtain, for every
2
m²
1cl12 S
Si = m
S
5
Summing over m I,.... T yields
T Σ m² (T -0--00 (fm))mm) m²
v ((2
11
1=
The remaining claims follow immediately from the standard equiva- lence between Hessian operator-norm bounds and local Lipschitz continuity of the gradient, and the reciprocal step-size rule 1/ < 1/L".
which proves equation 39
Theorem 4 (Relative smoothness and scaling the uniform mixture (restated)). At the uniform vector X = (1/m....,1) the loss s()) -(()) is Lm-shhoth relative to the KL-divergence, with
Thel S (T m) m
(40)
Consequently, the conservative step-size rule 1) S 1 Lmd yields the asymptotic scaling
1 θ -
(41)
fixed m.
Proof. Proposition Bestblishee that at X the Hessian satisfies
2 s(X)) s (T-m) (T-m) m
Relative smoothness of f with respect to the KL-divergence means V² < Liel V²y(A)) where (A) E, g Ag log X, is the negative entropy whose Bregman divergence is the KL. At the uniform vector, V ( = m Z (since m 2 1). Therefore
V² 3 (T () m)mm < m)-m) (')("")
which gives Ludi < (T- (T-m) m² as claimed. The conservative step-size choice S 1/ Lnel is therefore sufficient guarantee stability of the mirror-descent / exponentiated-grnee -radieen update. Since T m O(T',, the scaling 7 = (11/T) follows for fixed m.
The assumption equation 35 (strict positivity of every transition probability appearing in the likeli- hood) is the minimal condition that guarantees a uniform finite Lipschitz constant L over the entire simplex Am-1- If some transitions were zero, then for parameter vectors placing mass coor- dinates corresponding to zero transitions some denominators sA) could vanish and the Hessian operator would be unbounded (hence global L exists).
34

---

Published as a conference paper at ICLR 2026
J ASYMPTOTIC PROPERTIES OF THE LIKELIHOOD GRADIENT
In this section, we study the asymptotic properties of the gradient of the log- log-likelihood function. We are mostly interested in the asymptotic scaling of the gradient with the sequence length T'
Lemma 3 (Asymptotic Properties of the Gradient). Let the true parameter X (m) induce a Markov chain on the history space y that is aperiodic and irreducible on a finite state space. This implies the chain is geometrically ergodic. Let E denote expectation with respect to the stationary distribution of this chain. The mean of the score vector g(Y) exhibits the following asymptotic properties as Nobs + oo: The expected score vector scales linearly with Nohs T mn
go gx(X) = Ex-lg(Y)) = Notr + 0(1),
(42)
where v v(X*) is the constant vector of stationary expected single-step scores, whose h-th component is:
T(Yi-,, Yi)
[v(X*)]n mm. Eagu
(43)
プ ग (Y- Y.)
The (11) term represents a constant offset due initial conditions that does not grow with Nois-
Proof. Let first define the score contribution from single time step as the vector 2, (Y) E m
whose h-th component is given by:
(Y-h: h: Yi)
2 )=
Σj= π (Y ())), Yi)
The total score vector is the sum of these contributions over the observation period:
T
g Y = Σ YY),
where Nots T - m.
1=++1
The process (Zz)t>m is a sequence of random vectors. Because it is function of the underlying ergodic Markov chain (Y--m.... Yt), the sequence (2,) is also ergodic and its distribution con- verges to stationary distribution. By the linearity of expectation, the expected score is the sum of the individual expectations:
Ex (Y) = Σ Ex. i2 (Y)].
f==+11
The key assumption is the geometric ergodicity of the Markov chain. This implies that the distribu- tion of the state (Y-m,,.... Y,) converges exponentially fast to the unique stationary distribution, regardless of the initial state (Y1...., Y.). Consequently, the expectation Ex- [2,({] (Y)] converges exponentially fast to its vationary-state expectation, v(X*). This convergence can be quantified. There exist constant vector C and rate p (0, 1) such that for all t > m:
[Ex+(z((Y)] v(>*) < ||C||p"mm
We can now rewrite the sum of expectations:
Ex- (gg(
Σ (v(A*) + (x Z. (Y) v(*
1==+1
T
Σ v(2) +
Σ [==+1
(Ex- (Y)] (Y)]
[=
Nots v(**) + ET.
where Ey is the cumulative error term due to the process not having reached stationarity at early time steps. We can bound the norm of this error term:
T
Er = Σ (x -
v(x*))
S
Σ
EX- [2.]
v(*i)
Σ
|C
1=
t=++1
1=
35

---

Published as a conference paper at ICLR 2026
Let ls t The sum becomes a geometric series:
Erll S |CI
|C||
Σ
0²
p p
The sum of the error terms is bounded by finite constant that does not depend T or Nobs- Therefore, the error term is O(1). This proves the first claim: go()*) Nots v(X*) + 0 Note that (1) is also of (Noos), the term is asymptotically sub-linear.
36


Reviews:
[{"review_id": "YrP1T2ggLR", "reviewer": "Reviewer_9cwh", "summary": "The main problem the paper tackles is: How do Transformers perform in-context learning when the task requires inferring latent mixture weights on the fly instead of machine learning algorithms? To study this problem, the authors introduce a synthetic ICL task based on the Mixture of Transition Distributions (MTD) model. Methodologically, they explicitly construct a 3-layer disentangled Transformer that mirrors the one-step mirror descent algorithm. Empirically, the show that the trained transformer can approximate the theoretical construction.", "questions": "1. Why did the author choose the disentangled transformer architecture to analyze? What will happen if we add back the MLP layer?\n\n2. I think it would be helpful if the authors could elaborate on the choice of the MTD model as a tool to analyze the latent structure of the data. The authors already provide a formal introduction in Section 3, but it might strengthen the paper if that formal definition were explicitly linked back to the intuitive example in Figure 1, which clearly motivates why we want to analyze ICL beyond the machine-learning algorithms studied in previous works.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Extends the “transformers-as-optimizers” view from GD on fully observed tasks to mirror descent (exponentiated gradient) on latent mixture weights, a novel setting (latent variables over the simplex) rather than another linear-regression variant.\n* The paper is well-written and easy to understand. The proof idea illustration in the main text is very helpful.\n* The theoretical construction is backed up by the empirical results.", "weaknesses": "1. The theoretical analysis relies on a simplified, MLP-free Transformer architecture. Hence, it remains uncertain whether the same mirror-descent behavior holds in standard Transformers with MLPs (But I understand this is always the limitation of all constructive proof for almost all ICL theory papers)\n\n2. Proposition 3 and the constructed model require access to the transition matrix $\\pi*$, which limits applicability when base dynamics are not shared/known. This $\\pi*$ is later encoded in the output matrix $\\tilde{W}_O$. Doesn't this mean that the transformer has no chance to learn the transition matrix of the MTD model? If this is the case, how can the theoretical result justify that the transformer is \"learning the latent mixture model in context\" if it only learn the mixture weight, not the transition matrix? I'm not an expert in MTD, so I might misunderstand the concept, and will appreciate further explanation on this", "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": "laHGQL2Xbe", "reviewer": "Reviewer_cRmd", "summary": "This paper studies transformers trained on the MTD (Mixture of Transition Distributions) task. They show that a three layer disentangled transformer can naturally represent one step of mirror descent on the maximum likelihood objective. They also show that trained transformers have the same attention maps as their theoretical construction and obtain similar performance to mirror descent with optimal step sizes, supporting the hypothesis that transformers solve this task via mirror descent.", "questions": "- What do the weights of the trained disentangled transformer look like?\n- How would the performance of the disentangled transformer compare to a standard transformer on this task? Is it clear that the MLPs won't help?\n- For the multi-step transformer, were the step sizes for each step tuned independently or were they tied together? In the trained transformer can you read off the step sizes? Are they the same or does it use a schedule?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The task is a natural generalization of previous in-context learning tasks studied in the literature and removes the dependence on a fixed causal graph. The paper is also very well written, especially section 4 (the explicit construction). In addition, while the experiments are relatively limited, they do support the paper's thesis that transformers solve this ICL task via mirror descent.", "weaknesses": "While the paper provides a construction for a three layer disentangled transformer that can express one step of mirror descent, it doesn't study optimization so it's still unclear how gradient descent finds this solution. To support the claim that transformers learn this construction, the paper relies on experiments on disentangled transformers which show that both the attention maps and the overall performance match their explicit construction. However, there are a few weaknesses of the experiments:\n\n- They don't verify agreement in weight space between the trained transformer and their construction, they only verify the attention maps/KL divergence. Since the attention maps are input dependent, it's difficult to use a single attention map to verify equivalence of two architectures.\n- The agreement in attention maps isn't very convincing, for example while the first attention map matches the diagonal structure of the explicit construction, the actual weights appear to be quite far apart. Additionally, the middle attention map has a strong banded structure which is not captured by the explicit construction.\n- I believe that the paper only empirically validates their claim on the disentangled transformer. However, if they only verify final performance and attention maps, it would be good to include a \"standard\" transformer baseline with projections and MLPs, since you can still easily verify attention maps and performance on these. The main advantage of training a disentangled transformer is to look for weight-space agreement.\n- The only evidence that transformers perform multi-step mirror descent are the performance plots in Figure 5. It would be good to verify attention maps or ideally the weights themselves to check whether this is actually true or whether it's implementing a different estimator with similar performance (although this seems unlikely).\n\nMinor points:\n- In the first part of the abstract/paper, the claim is phrased as \"transformers implement Mirror Descent,\" however the real claim in these sections is that transformers *can* implement mirror descent, and then the empirical part of the paper tries to argue that they do.\n- line 151+152: While predicting $\\lambda$ is probably a necessary step for most good estimators, the actual task is just to predict $y_t$ given $y_{1:t-1}$, not to predict $\\lambda$?\n- line 227: overhang in the main paper", "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": "Axq3rqpIby", "reviewer": "Reviewer_vGtX", "summary": "This paper investigates the Mixture of Transition Distribution (MTD) task, where a latent variable must be learned in-context. The authors show that an ad hoc disentangled transformer is able to perform one step of Mirror Descent (MD), a mechanism derived from initializing on a specific prior distribution. This MD step serves as a first-order approximation of the Bayesian posterior mean. For longer sequences, they empirically demonstrate that Multi-step Mirror Descent approximates the Bayesian posterior mean more effectively. They also show that training disentangled transformers match the results obtained by their construction.", "questions": "1. On line 75 you claim the training of the transformer is using Gradient Descent, then you actually use ADAM and do not give a reason, I do not think they will both converge to the same critical point, could you comment on this?.\n2. Could you give a discussion where having a known transition matrix \\pi is realistic? I might not get why you were using rows from a Dirichlet distribution.\n3. How might non-linear activation functions (like ReLU or GELU) interplay with the smooth, exponentiated gradient algorithm required for the MD update?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The primary strength of this paper lies in its conceptual originality and technical rigor in tackling the challenge of dynamic latent variable inference. The paper's most original step is presenting the Mixture of Transition Distributions (MTD) model as an in-context learning task requiring the dynamic estimation of latent mixture weights. This approach successfully moves the field past simpler models (like linear regression or elementary Markov Chains). By identifying and proving that a disentangled transformer implements one step of the MD algorithm, the paper establishes a new insight: a trained disentangled transformer (or the attention part of transformers) in this context can effectively approximate the Bayesian posterior mean and the natural algorithm to achieve it. This provides an algorithmic explanation for how transformers update their beliefs about dynamic sequence properties, representing a good contribution to understanding transformer capacity.", "weaknesses": "I believe there might be a limitation in the scope of the generalization claim. To truly support the title's broad assertion that transformers learn via Mirror Descent, the paper should include an experiment demonstrating that this mechanism holds up in a standard transformer architecture (including MLPs). Furthermore, since the formal proof only covers a single MD step, it would be highly valuable to provide an explicit discussion or intuition regarding the theoretical path for two or more MD steps and whether an extra iteration naturally leads to a quadratic approximation of the Bayesian posterior mean.", "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": "sPg50Pbrhd", "reviewer": "Reviewer_wFmK", "summary": "This paper introduces a framework for in-context learning with latent variable modelling, based on a mixture of transition distributions. \nThe authors introduce a synthetic task that reframes the estimation of past tokens to learning latent mixture weights. The paper also proves that the learned algorithm is mirror descent. The paper also empirically shows that the learned algorithm for a **disentangled** transformer is mirror descent. Finally, the paper shows that the estimator is an approximation of the Bayesian optimal solution.", "questions": "Could more experiments be added for real tasks like language modelling?\nCould the models in experiments be extended to the general transformer?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper has the following strengths:\n1. Clear framework presentation with latent variable modelling.\n2. Good theory for the disentangled transformer\n3. Good link to Bayse optimal solution", "weaknesses": "1. While the theory is solid and I consider it a good contribution, the main problem is that the results are focused only on the disentangled transformer.\n - The paper claims that the results were proved for the general transformer (for example, in lines 480-481). \n-  The authors should make it clear that they refer to a disentangled transformer, which is not the same as the generally used one for practical tasks.\n2. The paper should try to extend results to the general case transformer, or state why they do not work. \n3. The experimental section should be improved.\n- Could more experiments be added for real tasks like language modelling?\n- Could the models in experiments be extended to the general transformer?\n4. Small typos:\n- 076, \"patterns math our\"? \n- 205-206: \"by a Transformer\" Do you mean a disentangled transformer?", "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": "GDHBiOwGBq", "reviewer": "Reviewer_y3Go", "summary": "The authors set up a synthetic task where to perform well at each step the transformer must estimate mixing weights for a set of transition distributions (the mixture of transition distributions task). They show that posterior mean inference for this task is intractable and then discuss one step of mirror descent as an attractive estimator of the mixing weights. They then show theoretically and empirically that disentangled transformers implement mirror descent to learn mixture weights for this process.", "questions": "- Is there any reason why 3 and 6 (or 5, there appears to be an inconsistency between the main text and the appendix) layers were chosen? Does it also hold for other layer numbers?\n\n- Are there any plans to extend this to a more realistic data generating process where tokens can depend on multiple previous tokens at different time lags?\n\n- Would the 5/6 layer disentangled transformer be amenable to the same sort of analysis as the three layer model, showing that it implements two step MD?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. I really like the matrix diagrams in section 4, it made following that section a lot easier. I felt that the paper was well written and structured. \n2. To my knowledge no one has previously proposed that ICL involves transformers performing mirror descent. This gives a new perspective on ICL\n3. I had not known of the MTD model before reading this, it appears to be a fairly well motivated synthetic task in this setting.", "weaknesses": "The lack of any experiments involving non-disentangled transformers weakens this papers claims. It's not clear to me that a general non-disentangled transformer would implement this mirror descent step. If you consider the Nichani et al. (2024) paper cited in this work, they also perform a theoretical and empirical analysis with disentangled transformers, but they include experiments involving standard non-disentangled transformers also. Something similar in this case would greatly improve the strength of this papers evidence for it's claims and make the title more justified.", "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)