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
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 ( π⋆ ) from dynamic, context-dependent weights ( λ ) 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., π⋆(dog,⋅) favouring verbs like fetch or chase ). Conversely, a low weight is assigned to the distractor (’bird’), suppressing its influence.
Figure 2: MTD for m=2.
Figure 3: Regularized Estimator. Comparison with Bayes and MLE estimators.
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 Softmax(W1A⊤) vs. true transitions matrix π⋆ 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 h^(l)
Concatenated Hidden State h(l)
0 (Input)
Token Embeddings
—
h(0)∈Rq
1
Responsibilities
h^(1)∈Rq+m
h(1)=\Concat(h(0),h^(1))∈R2q+m
2
Summation
h^(2)∈R2q+2m
h(2)=\Concat(h(1),h^(2))∈R4q+3m
3
Weighting
h^(3)∈R4q+3m
h(3)=\Concat(h(2),h^(3))∈R8q+6m
Final Output Matrix
WO∈Rq×(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 ( α )
Uniform ( αg=1 for all g )
Transition Matrix ( π )
Rows from Dirichlet ( α=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.
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 W1A⊤ ) 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.
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 .
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.
Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman +1
1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.
Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.
Chenyang Zhang, Yuan Cao
School of Computing & Data Science, The University of Hong Kong
Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that transformers can implement a gradient descent algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of attention depth in enabling the extrapolation abilities of transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.
Gyeonghun Kang, Changwoo J. Lee, Xiang Cheng
Department of Statistical Science, Duke University, Durham, NC, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA