Transformers Learn Latent Mixture Models In-Context via Mirror Descent

• URL: http://arxiv.org/abs/2604.10848v1
• Date: Sun, 12 Apr 2026 22:45:43 GMT
• Title: Transformers Learn Latent Mixture Models In-Context via Mirror Descent
• Abstract summary: Sequence modelling requires determining which past tokens are causally relevant from the context and their importance.<n>We introduce a framework based on Mixture of Transition Distributions, where a latent variable determines the influence of past tokens on the next.<n>We demonstrate that transformers can implement Mirror Descent to learn these weights from the context.
• Score: 36.758716760573215
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Induction Heads Interpolate N-Grams [29.939177622666858]
  Induction heads are attention circuits believed to underlie in-context learning in transformers.<n>We study transformers trained on order-$k$ Markov chains and identify two complementary smoothing mechanisms.<n>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.
  arXiv  Detail & Related papers  (2026-07-02T22:19:26Z)
• Incremental Learning of Sparse Attention Patterns in Transformers [29.54151079577767]
  This paper introduces a high-order Markov chain task to investigate how transformers learn to integrate information from multiple past positions.<n>We identify a shift in learning dynamics from competitive, where heads converge on the most statistically dominant pattern, to cooperative, where heads specialize in distinct patterns.
  arXiv  Detail & Related papers  (2026-02-22T12:16:06Z)
• Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration [52.017716672255524]
  Uncertainty calibration in pre-trained transformers is critical for their reliable deployment in risk-sensitive applications.<n>We propose a diffusion-inspired reconfiguration of transformers in which each feature transformation block is modeled as a probabilistic mapping.<n>Our method achieves superior calibration and predictive accuracy compared to existing uncertainty-aware transformers.
  arXiv  Detail & Related papers  (2026-02-09T17:24:47Z)
• Selective Induction Heads: How Transformers Select Causal Structures In Context [50.09964990342878]
  We introduce a novel framework that showcases transformers' ability to handle causal structures.<n>Our framework varies the causal structure through interleaved Markov chains with different lags while keeping the transition probabilities fixed.<n>This setting unveils the formation of Selective Induction Heads, a new circuit that endows transformers with the ability to select the correct causal structure in-context.
  arXiv  Detail & Related papers  (2025-09-09T23:13:41Z)
• On the Robustness of Transformers against Context Hijacking for Linear Classification [26.1838836907147]
  Transformer-based Large Language Models (LLMs) have demonstrated powerful in-context learning capabilities.<n>They can be disrupted by factually correct context, a phenomenon known as context hijacking.<n>We show that a well-trained deeper transformer can achieve higher robustness, which aligns with empirical observations.
  arXiv  Detail & Related papers  (2025-02-21T17:31:00Z)
• Dynamics of Transient Structure in In-Context Linear Regression Transformers [0.5242869847419834]
  We show that when transformers are trained on in-context linear regression tasks with intermediate task diversity, they behave like ridge regression before specializing to the tasks in their training distribution.<n>This transition from a general solution to a specialized solution is revealed by joint trajectory principal component analysis.<n>We empirically validate this explanation by measuring the model complexity of our transformers as defined by the local learning coefficient.
  arXiv  Detail & Related papers  (2025-01-29T16:32:14Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• A Theoretical Analysis of Self-Supervised Learning for Vision Transformers [66.08606211686339]
  Masked autoencoders (MAE) and contrastive learning (CL) capture different types of representations.<n>We study the training dynamics of one-layer softmax-based vision transformers (ViTs) on both MAE and CL objectives.
  arXiv  Detail & Related papers  (2024-03-04T17:24:03Z)
• In-Context Convergence of Transformers [63.04956160537308]
  We study the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent. For data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process.
  arXiv  Detail & Related papers  (2023-10-08T17:55:33Z)
• ExpPoint-MAE: Better interpretability and performance for self-supervised point cloud transformers [7.725095281624494]
  We evaluate the effectiveness of Masked Autoencoding as a pretraining scheme, and explore Momentum Contrast as an alternative. We observe that the transformer learns to attend to semantically meaningful regions, indicating that pretraining leads to a better understanding of the underlying geometry.
  arXiv  Detail & Related papers  (2023-06-19T09:38:21Z)
• Transformers learn in-context by gradient descent [58.24152335931036]
  Training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass.
  arXiv  Detail & Related papers  (2022-12-15T09:21:21Z)
• XAI for Transformers: Better Explanations through Conservative Propagation [60.67748036747221]
  We show that the gradient in a Transformer reflects the function only locally, and thus fails to reliably identify the contribution of input features to the prediction. Our proposal can be seen as a proper extension of the well-established LRP method to Transformers.
  arXiv  Detail & Related papers  (2022-02-15T10:47:11Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.