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

The factors driving the performance of in-context learning (ICL) in large language models (LLMs) remain poorly understood despite ICL's surprising effectiveness, enabling models to adapt to new tasks from only a handful of examples. To clarify and improve these capabilities, we characterize how the statistical properties of the pretraining distribution (e.g., tail behavior, coverage) shape ICL. We develop a theoretical framework that encompasses generalization and task selection and show how distributional properties govern sample efficiency, task retrieval, and robustness. To this end, we generalize existing concentration results to heavy-tailed priors and dependent sequences, better reflecting the structure of LLM pretraining data. Our framework reveals a fundamental design trade-off: heavy-tailed pretraining distributions facilitate robust task selection under distribution shifts but are detrimental to generalization, especially in low-data regimes. We then empirically evaluate our predictions by studying how ICL performance varies with the pretraining distribution on challenging tasks such as stochastic differential equations and stochastic processes with memory. Together, these findings suggest that controlling key statistical properties of the pretraining distribution is essential for building ICL-capable and reliable LLMs.

When we interact with a language model, we often give it a few examples of what we want before asking it to do something new. Remarkably, modern language models are very good at this — but we do not fully understand why some models are better than others. We studied how the variety of texts and tasks a model learns from during training shapes this ability. We discovered a surprising tension: giving a model a more varied training experience makes it more adaptable when facing unexpected situations, but also makes it much harder to train well. Surprisingly, this means that the same property that helps a model handle the unexpected actively hurts its ability to learn reliably — too much variety, and the model needs far more data to learn anything well. Too little, and it struggles the moment it encounters something new. This means there is no free lunch — building a model that is both reliable and adaptable requires carefully balancing the diversity of its training data, depending on how much data is available and how surprising the real-world tasks are expected to be.