Transformers Learn Latent Mixture Models In-Context via Mirror Descent

Francesco D'Angelo; Nicolas Flammarion

cs.LGApr 12, 2026

Transformers Learn Latent Mixture Models In-Context via Mirror Descent

Authors: Francesco D'Angelo, Nicolas Flammarion

Organizations: TML Lab, EPFL, Switzerland

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.

Figures & tables

Panel of Figure 1 — Figure 1: To predict the final word, the model must infer the causal relevance of past tokens. Our MTD framework models this by separating a static, context-free unigram ( $\pi^{\star}$ ) from dynamic, context-dependent weights ( $\bm{\lambda}$ ) that are inferred in-context. The model learns to assign high weights to the causally relevant positions (’dog’) and (’ball’), activating their respective slices of the unigram (e.g., $\pi^{\star}(\text{dog},\cdot)$ favouring verbs like fetch or chase ). Conversely, a low weight is assigned to the distractor (’bird’), suppressing its influence.

Panel of Figure 2 — Figure 2: MTD for m=2.

Figure 3: Regularized Estimator. Comparison with Bayes and MLE estimators.

Panel of Figure 4 — 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 to the ground truth transition probabilities for the trained transformers, the constructed transformer, and the one-step MD estimator across sequence lengths. Right (bottom): First-layer attention softmax $\text{Softmax}(\bm{W}_{1}^{A\top})$ vs. true transitions matrix $\pi^{\star}$ for a trained model.

Figure 5: Multi-Step MD vs. 5-Layer Transformer. KL divergence to the ground-truth transition probabilities for a 5-layer trained transformer, the

k

-step MD estimators and the Bayes-optimal estimator across sequence lengths.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Layer	Description	Attention Output $\hat{\bm{h}}^{(l)}$	Concatenated Hidden State $\bm{h}^{(l)}$
0 (Input)	Token Embeddings	—	$\bm{h}^{(0)}\in\mathbb{R}^{q}$
1	Responsibilities	$\hat{\bm{h}}^{(1)}\in\mathbb{R}^{q+m}$	$\bm{h}^{(1)}=\Concat(\bm{h}^{(0)},\hat{\bm{h}}^{(1)})\in\mathbb{R}^{2q+m}$
2	Summation	$\hat{\bm{h}}^{(2)}\in\mathbb{R}^{2q+2m}$	$\bm{h}^{(2)}=\Concat(\bm{h}^{(1)},\hat{\bm{h}}^{(2)})\in\mathbb{R}^{4q+3m}$
3	Weighting	$\hat{\bm{h}}^{(3)}\in\mathbb{R}^{4q+3m}$	$\bm{h}^{(3)}=\Concat(\bm{h}^{(2)},\hat{\bm{h}}^{(3)})\in\mathbb{R}^{8q+6m}$
Final Output Matrix		$\widetilde{\bm{W}}_{O}\in\mathbb{R}^{q\times(8q+6m)}$

Appendix

Table 1: Dimensionality of Hidden States in the Disentangled Transformer

Component	Parameter	Value
Data Generation
	MTD model order ( $m$ )	3,4,5
	Vocabulary size ( $q$ )	5
	Sequence Length ( $T$ )	Varied (64 to 1984)
	Dirichlet Prior ( $\bm{\alpha}$ )	Uniform ( $\alpha_{g}=1$ for all $g$ )
	Transition Matrix ( $\bm{\pi}$ )	Rows from Dirichlet ( $\bm{\alpha}=1$ )

Appendix

Table 2: Summary of Experimental Parameters

Figure 6: KL divergence to the ground-truth We report the KL 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.

Panel of Figure 7 — Figure 7: Layer-1 Attention Softmax vs. True Transition matrices for orders 3 and 5. For both orders ( $m=3$ top 3 rows, $m=5$ bottom 3 rows), each panel reports the learned first-layer attention (softmax of $\bm{W}_{1}^{A\top}$ ) alongside the true transition matrix; the average row-wise KL divergence is reported in the panel title. Sequence lengths increase left-to-right, top-to-bottom.

Panel of Figure 8 — Figure 8: Comparison of Trained and Constructed Transformers (attention grids). Top: Attention maps of the trained transformer (disentangled) versus our theoretical construction (seq. length 64, MTD order $m=3$ ). Bottom: Same as top but for MTD order $m=5$ .

Explore similar work

Induction Heads Interpolate N-Grams
Jul 2, 2026Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman +1Transformer Interpretability In-Context Learning
1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.
Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent
May 7, 2026Chenyang Zhang, Yuan CaoLogistic Regression Self-Attention
School of Computing & Data Science, The University of Hong Kong
Transformers Can Learn Posterior Predictive Distributions In-Context
May 26, 2026Gyeonghun Kang, Changwoo J. Lee, Xiang ChengBayesian Inference Transformer Attention
Department of Statistical Science, Duke University, Durham, NC, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA

Jul 2, 2026J/K move · Enter open · S save

Induction Heads Interpolate N-Grams

Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman, Nicolas Flammarion

1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.

Transformer Interpretability In-Context Learning Language Modeling Induction Heads

Open paper AskPDF arXiv

Induction heads are attention circuits believed to underlie in-context learning in transformers, yet a precise characterization of the estimators they implement remains elusive. We study transformers trained on order-

k

Markov chains and identify two complementary smoothing mechanisms. First, at finite attention-weight scale, the circuit implements a soft context-matching estimator: it aggregates contributions from exact and partial context matches, weighted exponentially by their overlap, and induces a data-dependent interpolation across context orders analogous to Jelinek-Mercer smoothing. Second, a beginning-of-sequence (BOS) token induces additive pseudo-counts, recovering Dirichlet-style smoothing. We construct a disentangled transformer implementing both mechanisms and show that trained transformers recover the predicted attention patterns. Across settings where pseudo-count smoothing is optimal or lower-order contexts provide structured evidence, trained transformers match or outperform classical count-based baselines. Our results bridge mechanistic interpretability of induction heads with classical statistical smoothing, revealing that transformers learn to regularize in-context estimation rather than simply count.

Abstract

Explore similar work

Induction Heads Interpolate N-Grams

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent

Transformers Can Learn Posterior Predictive Distributions In-Context