A Unified Perspective on Task Retrieval and Learning in In-Context Learning on Markov Data
Abstract
In-context learning enables pretrained transformers to adapt from a prompt without parameter updates, but it remains unclear when this behavior arises from retrieving memorized pretraining tasks (i.e., task retrieval), versus learning a rule from in-context data (i.e., task learning). We study this distinction through a finite-mixture Markov data model, where each latent Markov chain represents a pretraining task. We characterize the finite-prior Bayes predictor induced by next-token pretraining and use it to explain when pretrained transformers exhibit retrieval or learning. In low-diversity regimes, this predictor retrieves the Markov chain for pretraining most compatible with the prompt, with error decaying exponentially in prompt length. This mechanism is realizable by a two-layer single-head causal transformer. In high-diversity regimes, the finite-prior Bayes predictor approaches the empirical Markov predictor inferred from the prompt, and architectural constraints can bias globally pretrained transformers toward this capacity-efficient task-learning mechanism, with error decreasing as prompt length increases. Finally, we extend our analysis to multi-state higher-order Markov chains.