Title: Transformers Learn Latent Mixture Models In-Context via Mirror Descent
Primary area: foundation or frontier models, including LLMs
Keywords: in-context learning, markov chain, transformers, mirror descent, mixture models, latent variables

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 novel framework based on Mixture of Transition Distributions, whereby a latent variable, whose distribution is parameterized by a set of unobserved mixture weights, determines the influence of past tokens on the next. To correctly predict the next token, transformers need to learn the mixture weights in-context. We first demonstrate that transformers implement Mirror Descent to learn the mixture weights from the context. We then provide a construction of a three-layer transformer that exactly implements one step of Mirror Descent and establish its validity by showing it yields a first-order approximation of the Bayes-optimal solution. Finally, we empirically show that transformers trained from scratch on this task converge to the constructed solution, with their learned attention patterns matching our construction.