📊 statistics

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

This paper establishes a theoretical and empirical framework demonstrating that the statistical properties of pretraining data, particularly heavy-tailed distributions, create a fundamental trade-off where robust task selection under distribution shifts comes at the cost of generalization performance in low-data regimes.

Original authors: Waïss Azizian, Ali Hasan

Published 2026-06-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Waïss Azizian, Ali Hasan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a large language model (like a super-smart AI) as a chef who has spent years cooking in a massive, chaotic kitchen. This chef hasn't been taught specific recipes for every single dish in the world. Instead, they have tasted thousands of different meals, from simple soups to complex stews, and learned the general "vibe" of cooking.

In-Context Learning (ICL) is like this chef being handed a new, strange recipe with just three examples of how to make it, and then immediately cooking it perfectly without ever having seen that specific dish before.

This paper asks: What kind of "kitchen" (pretraining data) makes the chef best at this trick?

The authors discovered a fundamental trade-off (a "pick two, lose one" situation) based on the statistical "shape" of the data the chef learned from.

1. The Two Types of "Kitchens" (Pretraining Distributions)

The paper compares two main types of data distributions the chef could have learned from:

The "Safe" Kitchen (Light-Tailed): Imagine a kitchen where almost every ingredient is common and predictable. You mostly see potatoes, carrots, and chicken. Extreme ingredients (like a rare, exotic spice) are almost non-existent. The data follows a neat, bell-curve pattern.
The "Wild" Kitchen (Heavy-Tailed): Imagine a kitchen where, yes, you see lots of potatoes, but you also frequently encounter wild, rare, and extreme ingredients. The "tail" of the distribution is long, meaning the chef has seen a huge variety of weird, outlier tasks.

2. The Great Trade-Off

The paper reveals that the choice of kitchen creates a tug-of-war between two skills:

Skill A: Spotting the New Task (Task Selection)

The Wild Kitchen Wins: If the chef trained in the "Wild" kitchen (heavy-tailed data), they are excellent at looking at a new, strange recipe and saying, "Ah, I've seen something like this before!" They are very robust. Even if the new task is weird or far outside the norm, they can adapt quickly because their training data included many "weird" examples.
The Safe Kitchen Loses: If the chef only trained on common ingredients, they get confused by the weird stuff. They struggle to identify what the new task is if it's too different from their training.

Skill B: Mastering the Basics (Generalization)

The Safe Kitchen Wins: If the chef trained in the "Safe" kitchen, they are incredibly precise at cooking the common dishes. When the new task is similar to what they've seen, they generalize perfectly with very few examples.
The Wild Kitchen Loses: Here is the catch. Because the "Wild" kitchen was so chaotic and full of rare outliers, the chef is a bit "scattered." If you give them a new task that is actually quite common, they might need many more examples to figure it out compared to the "Safe" chef. The heavy-tailed data makes it harder for them to lock onto the standard patterns when data is scarce.

3. The "Long Memory" Problem

The paper also looked at dependencies. Imagine the chef is cooking a stew where the taste of the current spoonful depends heavily on what was in the pot 10 minutes ago (long-range memory).

The Finding: If the data has strong, long-range dependencies (like a complex, evolving story or a process with memory), the chef needs even more training tasks to learn well.
The Analogy: It's like trying to learn a language where the meaning of a word today depends on a word you read three chapters ago. If the "kitchen" is full of these complex, long-winding stories, the chef needs to taste thousands more recipes to get the hang of it, especially if the data is "wild" (heavy-tailed).

4. The Bottom Line

The paper concludes that there is no perfect "one-size-fits-all" training diet for these models.

If you want your AI to be resilient and adaptable to weird, new, or shifting situations (like a robot in a disaster zone), you should feed it heavy-tailed data (lots of variety, including outliers). But be prepared: it might need more examples to learn a standard task.
If you want your AI to be highly efficient at learning standard tasks with very few examples, you should feed it light-tailed data (more consistent, predictable patterns). But it will likely fail if the task is too strange or different from its training.

In short: You can't have the best of both worlds. The "Wild" kitchen makes the chef a great detective for new mysteries but a slower learner for standard recipes. The "Safe" kitchen makes the chef a master of standard recipes but a poor detective for the unknown. The paper provides the mathematical proof for this balance.

Technical Summary: How Does the Pretraining Distribution Shape In-Context Learning?

Problem Statement

In-context learning (ICL) allows large language models (LLMs) to adapt to new tasks using only a few examples provided in the input context, without updating model weights. Despite its effectiveness, the factors driving ICL performance remain poorly understood. While previous works have explored ICL through algorithmic, Bayesian, or empirical lenses, there is a lack of a principled framework connecting the statistical properties of the pretraining distribution (e.g., tail behavior, coverage, temporal dependence) to ICL behavior at test time. Specifically, existing literature often assumes light-tailed (e.g., Gaussian) priors and independent and identically distributed (i.i.d.) data, failing to capture the heavy-tailed and dependent structures characteristic of real-world LLM pretraining data. This paper addresses the gap in understanding how these distributional properties govern sample efficiency, task retrieval, and robustness under distribution shifts.

Methodology

The authors develop a theoretical framework that decomposes ICL performance into two distinct components: task selection (identifying the correct task from the context) and generalization (performing well on unseen tasks).

Theoretical Framework:
- Modeling: Training data is modeled as a mixture of tasks sampled from a prior distribution $\pi(\theta)$ . Data sequences within tasks are generated according to task-specific distributions that may be dependent and non-Markovian.
- Heavy-Tailed & Dependent Analysis: The framework generalizes existing concentration results to handle heavy-tailed priors (specifically Student- $t$ and Generalized Normal distributions) and dependent sequences (using weak dependence assumptions via Wasserstein distance).
- Generalization Error: The authors derive concentration bounds for the generalization error of a trained model minimizing next-sample prediction loss. They establish conditions under which the model enters the "ICL regime" (approximating the Bayes-optimal predictor for the true prior) rather than the "in-weight learning" regime.
- Task Selection: Adopting a Bayesian perspective, the authors analyze the concentration of the posterior distribution over tasks given in-context examples. They quantify the speed at which the posterior concentrates around the true task $\theta^*$ , linking this rate to the prior density $\pi(\theta^*)$ .
Empirical Validation:
- The theoretical predictions are validated on challenging numerical tasks: Linear Regression, Ornstein-Uhlenbeck (OU) processes (stochastic differential equations), and Stochastic Volterra equations (processes with long-range memory).
- Experiments vary the tail heaviness of the pretraining distribution (via degrees of freedom $\nu$ for Student- $t$ and shape parameter $\beta$ for Generalized Normal) and the number of pretraining tasks ( $N$ ).
- Performance is evaluated via Mean Squared Error (MSE) under varying magnitudes of distribution shift (task center displacement).

Key Contributions

A General Theoretical Framework: The paper provides a unified theoretical analysis of ICL that accommodates heavy-tailed priors and non-i.i.d. dependent sequences, reflecting the structure of LLM pretraining data more accurately than prior Gaussian/i.i.d. assumptions.
Identification of a Fundamental Trade-Off: The core theoretical contribution is the revelation of a design trade-off in pretraining distribution selection:
- Heavy-tailed priors facilitate robust task selection and better performance under distribution shifts (out-of-distribution tasks) because they assign non-negligible probability mass to tasks far from the prior's mode.
- However, heavy-tailed priors are detrimental to generalization in low-data regimes (small $N$ ), as they require significantly more pretraining tasks to achieve the same level of generalization error compared to light-tailed priors.
Quantification of Temporal Dependence: The framework demonstrates that stronger temporal dependence in data sequences (e.g., long-range memory in Volterra processes) negatively impacts generalization, increasing the number of pretraining tasks required for reliable ICL.
Empirical Validation: The authors empirically confirm that heavier tails improve robustness to distribution shifts but degrade performance when the number of pretraining tasks is limited, validating the predicted trade-off across linear regression and stochastic processes.

Key Results

Generalization Bounds (Theorem 1): The generalization error depends critically on the maximal moment $q$ of the prior $\pi$ . For heavy-tailed priors (small $q$ ), the error bound worsens, particularly for small deviation probabilities ( $\delta$ ). The number of tasks $N$ required for good generalization scales with the tail heaviness and the strength of temporal dependence ( $A_T, B_T$ ).
Task Selection Bounds (Theorem 2): The convergence rate of the posterior distribution to the true task is governed by $\log(1/\pi(\theta^*))$ . For tasks far from the bulk of the prior, heavy-tailed distributions (where $\pi(\theta^*)$ decays slowly) result in faster convergence and lower error compared to light-tailed distributions.
Experimental Findings:
- Linear Regression: Heavy-tailed priors ( $\nu=3$ ) outperform light-tailed priors ( $\nu=\infty$ ) under large distribution shifts but perform worse when the number of pretraining tasks ( $N$ ) is small (e.g., $N=100$ ). As $N$ increases, the heavy-tailed advantage in robustness emerges.
- Stochastic Processes: Similar trends are observed in OU and Volterra processes. Stronger temporal dependencies (smaller kernel exponent $\alpha$ in Volterra) exacerbate the generalization penalty, requiring more tasks to achieve comparable performance.

Significance and Claims

The paper claims to provide a principled way to translate pretraining distributional properties into test-time ICL behavior. It argues that controlling key statistical properties of the pretraining distribution is essential for building ICL-capable and reliable transformers.

Practical Implications: The findings suggest that for domains requiring robustness to distribution shifts (e.g., robotics, healthcare), a pretraining distribution with heavier tails (long-tailed mixture over latent tasks) may be beneficial, provided sufficient pretraining data is available to mitigate generalization penalties. Conversely, for tasks with limited pretraining data, lighter tails may be preferable to ensure generalization.
Theoretical Advancement: By extending concentration inequalities to dependent, heavy-tailed sequences, the work bridges the gap between theoretical learning guarantees and the empirical realities of LLM pretraining, offering a new lens to understand the emergence and limitations of ICL.

The authors remain modest regarding their scope, noting that their experiments are limited to controlled numerical tasks and that their theoretical guarantees assume the model is sufficiently expressive to approximate the Bayes-optimal predictor. They identify extending these insights to large-scale natural language datasets as a critical direction for future work.