How Does the Pretraining Distribution Shape In-Context Learning? A Fundamental Trade-Off

In-context learning (ICL) has emerged as a central capability of pretrained language models, yet its theoretical analysis has focused primarily on causal language models trained by left-to-right autoregressive prediction, such as GPT-style models. Masked language models instead recover masked tokens from bidirectional context, and their role in ICL remains less understood. We develop a statistical learning framework that represents the context examples by their empirical measure and models prediction as a function of the context and the query. This formulation places autoregressive and masked pretraining objectives within a common excess-risk analysis. Under Wasserstein-type regularity conditions, we relate pretraining with T tasks and N samples per task to k-shot excess risk at inference, obtaining same-order upper bounds for masked and autoregressive objectives. We also study task-distribution shift, where pretraining tasks are sampled from P and inference tasks from Q; the resulting bound contains an additional term controlled by the lifted Wasserstein distance between P and Q. The bounds further imply an order-optimal allocation under a fixed pretraining data budget and refined rates under intrinsic low-dimensional structure. Experiments on controlled function-learning tasks show that the Masked Pair Encoder (MPE) can achieve performance comparable to GPT-2-style causal Transformers, suggesting that ICL behavior is not specific to causal language models.

In-context learning (ICL) is an emerging paradigm that employs the semantic information inherent in large language models (LLMs) for generating answers to user queries. While the remarkable performance of ICL has been widely known, a general modeling and a rigorous theoretical analysis of this paradigm are still lacking. This work presents a probabilistic model for ICL and derives the performance of ICL for both general parametric distributions and exponential families. Based on the derived results, the work explains the impact of multiple factors such as the number of demonstrations, the sensitivity of the probabilistic model to the variation of its parameters, as well as the similarity between the demonstrations and the query on the performance of ICL.

Large language models (LLMs) are increasingly used as conditional generators for structured data, relying on in-context learning (ICL) to adapt to new distributions without parameter updates. We investigate the limits of ICL for structured generation under distribution mismatch, using high-cardinality tabular data as a controlled test case, and identify a structural failure mode we term \textit{categorical prior lock-in}: the inability of ICL to update the model's prior over token distributions inherited from pre-training. Across two 7B-parameter open-weight models, ICL improves numerical fidelity with additional examples but exhibits a sharp ceiling on categorical distributions, failing to reproduce rare classes entirely. Parameter-efficient fine-tuning (LoRA) overcomes these limitations but introduces measurable memorization risk and, in some cases, destabilizes structured output generation, highlighting a fundamental trade-off between adaptability and privacy.