Transformers Learn Latent Mixture Models In-Context via Mirror Descent

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.

• The paper demonstrates that transformers implement a one-step mirror descent update to infer latent mixture weights, effectively bridging attention mechanisms with statistical estimation.
• The study formalizes the Mixture of Transition Distributions framework, detailing how transformer layers compute soft assignments and aggregate evidence from past tokens.
• Empirical validations confirm that both three-layer and deeper transformers accurately match the theoretical mirror descent predictions, with attention patterns mirroring true transition matrices.

Latent Mixture Model Inference in Transformers via Mirror Descent

Introduction and Motivation

The work "Transformers Learn Latent Mixture Models In-Context via Mirror Descent" (2604.10848) provides a theoretically rigorous and empirically validated analysis of the mechanisms by which transformers perform statistical inference over latent mixture models directly in-context. The central focus is on the Mixture of Transition Distributions (MTD), which generalizes higher-order Markov chains through latent mixture weights, capturing more flexible causal dependencies than standard Markov approaches. This goes beyond prior work, which addressed only settings with static dependencies and no structural latent variables.

The main contribution is the identification, formalization, and explicit implementation (constructive proof) of mirror descent as the algorithmic procedure underlying transformers' inference of latent mixture weights from context, both in theoretical construction and in gradient-trained models. This interpretation expands the gradient-based view of in-context learning and demonstrates transformers' ability to dynamically adapt causal structure estimation directly from sequence context, a major step in connecting mechanistic interpretability to statistical estimators for practical sequence modeling.

MTD Framework and In-Context Learning Task

The paper formalizes the sequence modeling task as estimating the relevance (importance) of past tokens as latent mixture weights $\bm{\lambda}$ in the MTD model. Each latent variable $Z_t$ selects a lag (position in the previous $m$ tokens), and a transition probability is drawn according to a common transition matrix $\bm{\pi}$. Sequence-specific $\bm{\lambda}$ vectors are sampled from a Dirichlet prior; the model must infer these mixture weights solely from a single observed sequence. This encapsulates the core in-context challenge: learning to infer and reweight latent causal contributions from observed data rather than from static correlation statistics.

The Bayes-optimal solution would entail integrating over the (non-conjugate) posterior of $\bm{\lambda}$, an intractable operation. Maximum-likelihood or MAP estimators are also analytically intractable, thus necessitating iterative approaches like EM or, as advocated here, mirror descent (MD).

Mirror Descent Approximation and Transformer Construction

The theoretically central result is that a three-layer transformer can exactly implement a one-step exponentiated gradient (mirror descent) update on the log-likelihood of the mixture parameters, initialized at the simplex barycenter (uniform prior). The explicit construction employs:

• Layer 1: Attention computes posterior responsibilities (soft assignments $\gamma_k(g)$ for each lag $g$ at each position $k$), leveraging carefully designed content-based attention and relative positional encoding.
• Layer 2: Aggregates (sums) lag-specific responsibilities across the sequence segment through uniform positional attention, storing results in the residual stream.
• Layer 3: Reweights and normalizes these summary statistics via attention and a trainable learning-rate scaling parameter $\beta$, producing the estimated mixture weights through a softmax.
• Output Layer: Applies the transition matrix $Z_t$0 to a $Z_t$1-weighted sum of one-hot encodings to recover the next-token predictive distribution.

The construction is made fully explicit and algebraically transparent, with each layer's function and parameterization derived from the mechanics of the MD algorithm. This yields a direct equivalence (for 1-step MD initialized at the barycenter) between the transformer's output and a regularized, first-order approximation of the Bayes predictor.

Theoretical Justifications and Algorithmic Optimality

A first-order Taylor analysis demonstrates that the one-step MD estimator and the true Bayesian posterior mean are locally equivalent for learning rate $Z_t$2, i.e., in the regime of no evidence (low data). The required step size $Z_t$3 is shown to scale as $Z_t$4 for sequence length $Z_t$5, matching the transformer implementation via the scaling $Z_t$6. This scaling is justified by bounding the loss function's Lipschitz constant (relative smoothness) and the Hessian norm uniformly over the simplex.

The regularizing properties of early-stopped MD are also articulated: iterated MD converges to the (suboptimal) MLE, but the early stopping point aligns much more closely to the (entropy-)regularized Bayesian mean, consistent with observed transformer behavior.

Empirical Validation

Empirical evaluation confirms the theory: both three-layer and deeper (five-layer) disentangled and standard transformers, trained from scratch, learn predictive distributions and attention patterns that closely match the theoretical construction and the mirror descent estimator across various sequence lengths, MTD orders ($Z_t$7), and initialization schemes.

Figure 2: KL divergence and attention patterns comparing trained (disentangled/standard) and constructed transformers with the one-step MD estimator, all closely align across sequence lengths.

• Trained models nearly exactly match the theoretical construction in terms of attention, output KL divergence to ground-truth transition probabilities, and the match between the first-layer attention softmax and the transition matrix.
• For increased depth, transformers approximate multi-step MD estimators (e.g., a five-layer transformer tracks two-step MD performance), implying that additional layers can compose multiple inference update steps, mimicking iterative regularized optimization.

Figure 1: KL divergence to ground-truth as a function of sequence length, validating the match between multi-step MD estimators and transformer predictions at order $Z_t$8 and $Z_t$9.

Figure 3: Learned first-layer attention weights (softmax of $m$0) closely mirror the rows of the true transition matrix for both $m$1 and $m$2.

Figure 4: Attention grids for trained and constructed transformers at $m$3 and $m$4 confirm banded lower-triangular structure and accurate computation of responsibilities per theoretical predictions.

Implications and Future Directions

This work concretely establishes gradient-based optimization (specifically, an MD step) as the precise mechanism through which transformers perform in-context inference of latent mixture weights under MTD sequence models. This closes the interpretability circuit for how transformers can implement latent variable inference, not just direct parameter regression/counting, and provides an explicit algorithmic circuit for token-importance discovery—central for capturing higher-order dependencies and dynamic causal structure in real language.

Key theoretical and practical implications include:

• Algorithm Selection and Circuit Decomposition: The analysis clarifies how transformers can perform task-adaptive algorithm learning, reweighting in-context evidence to infer which generative rules or memory traces to attend to, not merely regressing static mappings.
• Mechanistic Interpretability: By providing a full constructive proof, the work bridges the gap between phenomenological circuit discovery in pre-trained LMs and algorithmic regularization principles, yielding interpretable network architectures designed to reflect statistical estimation algorithms.
• Generalization to Other Latent Structures: The approach can, in principle, be generalized to richer latent variable models (e.g., latent class HMMs or probabilistic context-free grammars), where tractable approximation and in-context inference are critical yet mechanistically uncharacterized.
• Design of Pretraining Objectives: The results suggest that transformer architectures (even when strongly parameter-constrained, e.g., no MLPs and explicit residual streams) can acquire dynamic, interpretable algorithmic skills, informing the construction of synthetic pretraining tasks for probing hierarchy and abstraction in sequence models.

Further research can extend this approach to multi-step MD circuits, latent-variable models with non-Dirichlet priors or nonparametric mixture spaces, and cross-modal sequential inference scenarios.

Conclusion

This paper gives a concrete, algorithmically explicit account of in-context learning for latent mixture models in transformers, establishing that suitably constructed (or gradient-trained) transformers realize mirror descent inference over mixture weights, achieving first-order Bayes-optimal performance both in theory and empirical practice. The connection provides a novel algorithmic lens for understanding attention-based in-context model selection and compositionality, as well as new foundations for sequential model interpretability and architecture design.