Markov Chain Estimation with In-Context Learning

TL;DR

This paper explores how transformer models can learn to estimate transition probabilities in Markov chains through in-context learning, demonstrating that larger models and diverse encodings improve generalization beyond memorization.

Contribution

It shows that transformers can learn to estimate Markov chain transition matrices from context, with thresholds in size and training data enabling generalization to unseen matrices.

Findings

Transformers can learn to predict transition probabilities in Markov chains.

Model size and training set size are critical for successful in-context learning.

Enhanced state encoding improves robustness to different Markov chain structures.

Abstract

We investigate the capacity of transformers to learn algorithms involving their context while solely being trained using next token prediction. We set up Markov chains with random transition matrices and we train transformers to predict the next token. Matrices used during training and test are different and we show that there is a threshold in transformer size and in training set size above which the model is able to learn to estimate the transition probabilities from its context instead of memorizing the training patterns. Additionally, we show that more involved encoding of the states enables more robust prediction for Markov chains with structures different than those seen during training.