Markov Chain Estimation with In-Context Learning
This addresses the problem of understanding in-context learning mechanisms in transformers for researchers in machine learning, though it is incremental as it builds on existing work in transformer capabilities.
The study explored whether transformers can learn to estimate Markov chain transition probabilities from context using next-token prediction, finding that beyond certain size and training data thresholds, models succeed without memorizing patterns, with improved encoding enhancing robustness to unseen chain structures.
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.