Transformers Learn Latent Mixture Models In-Context via Mirror Descent

TL;DR

This paper formalizes how transformers learn to estimate token importance in context by implementing Mirror Descent, providing theoretical constructions and empirical evidence of this mechanism.

Contribution

It introduces a novel framework linking transformers to latent mixture models and demonstrates how they can implement Mirror Descent to learn token relevance.

Findings

Transformers can implement Mirror Descent to learn token importance.

A three-layer transformer can exactly perform one step of Mirror Descent.

Empirical results show trained transformers align with the theoretical model.

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.