Title: Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression
Venue: NeurIPS 2023
Accepted: yes
Keywords: in-context learning

Bayesian inference

transformers

task diversity

emergence
Abstract: Pretrained transformers exhibit the remarkable ability of in-context learning (ICL): they can learn tasks from just a few examples provided in the prompt without updating any weights. This raises a foundational question: can ICL solve fundamentally _new_ tasks that are very different from those seen during pretraining? To probe this question, we examine ICL’s performance on linear regression while varying the diversity of tasks in the pretraining dataset. We empirically demonstrate a _task diversity threshold_ for the emergence of ICL. Below this threshold, the pretrained transformer cannot solve unseen regression tasks, instead behaving like a Bayesian estimator with the _non-diverse pretraining task distribution_ as the prior. Beyond this threshold, the transformer significantly outperforms this estimator; its behavior aligns with that of ridge regression, corresponding to a Gaussian prior over _all tasks_, including those not seen during pretraining. Thus, when pretrained on data with task diversity greater than the threshold, transformers _can_ optimally solve fundamentally new tasks in-context. Importantly, this capability hinges on it deviating from the Bayes optimal estimator with the pretraining distribution as the prior. This study also explores the effect of regularization, model capacity and task structure and underscores, in a concrete example, the critical role of task diversity, alongside data and model scale, in the emergence of ICL.

Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression

Abstract

Pretrained transformers exhibit the remarkable ability of in-context learning (ICL): they can learn tasks from just a few examples provided in the prompt without updating any weights. This raises a foundational question: can ICL solve fundamentally new tasks that are very different from those seen during pretraining? To probe this question, we examine ICL's performance on linear regression while varying the diversity of tasks in the pretraining dataset. We empirically demonstrate a task diversity threshold for the emergence of ICL. Below this threshold, the pretrained transformer cannot solve unseen regression tasks, instead behaving like a Bayesian estimator with the non-diverse pretraining task distribution as the prior. Beyond this threshold, the transformer significantly outperforms this estimator; its behavior aligns with that of ridge regression, corresponding to a Gaussian prior over all tasks , including those not seen during pretraining. Thus, when pretrained on data with task diversity greater than the threshold, transformers can optimally solve fundamentally new tasks in-context. Importantly, this capability hinges on it deviating from the Bayes optimal estimator with the pretraining distribution as the prior. This study also explores the effect of regularization, model capacity and task structure and underscores, in a concrete example, the critical role of task diversity, alongside data and model scale, in the emergence of ICL.

1 Introduction

Pretrained transformers (PTs) can learn new tasks from just a few examples provided in the prompt without taking any gradient steps on those examples [1]. This ability, called in-context learning (ICL) , has unlocked the widspread use of language models by making it efficient to adapt general purpose models to bespoke tasks without explicit training. Though remarkable, what makes ICL mysterious, and potentially harmful [2], is that the learning algorithm implemented by the PT in its forward pass is not built into its architecture or training process; instead it emerges from pretraining on large-scale data with a next token prediction objective. This raises a foundational question: can ICL really solve fundamentally new tasks that are very different from those seen during pretraining? If so, what learning algorithm does ICL implement? To answer these questions, we need to better understand how the different ingredients that go into pretraining influence this ability.

Towards this end, we explore how the diversity of tasks in the pretraining data affects the emergence of ICL. Prior work [3] has proposed that ICL works by performing Bayesian inference. During pretraining, transformers learn a prior over latent tasks represented in the pretraining data. When prompted with examples at inference time, they "retrieve" relevant pretraining tasks and generate subsequent tokens from the posterior distribution conditioned on the query and inferred tasks. This suggests that ICL performance on a new task is influenced by its similarity to tasks implicitly learned during pretraining. However, the distribution of tasks in our pretraining data, T Pretrain, is usually a

limited and unrepresentative subsample of the ideal distribution of tasks, T True, that we want our model to be capable of learning in-context. For instance, T True could be the set of all instructions we want an A.I. assistant to follow. But, large-scale language modeling datasets [4, 5] used to pretrain these models contain very few examples correctly formatted for ICL. Instruction finetuning (IFT) datasets [6-11] designed to ameliorate this are expensive to collect and thus contain tasks from just a few domains. Under the Bayesian framework, this distribution mismatch would cause the Bayesian estimator with a prior over the limited pretraining tasks, T Pretrain, to perform suboptimally on tasks that are very different from those seen during pretraining. This motivates our question: can a model pretrained on a dataset with low task diversity nevertheless learn new, unseen tasks?

For general purpose language modeling, the size and complexity of T Pretrain and the vague specification of T True make this question challenging to analyze. So, following recent work [12-14], we study ICL for linear regression. Here, a task is a linear regression problem with a given latent regression vector; the PT must predict the target for a new data point from examples of data-target pairs provided in the prompt. Prior work [13] has shown that transformers that see an unlimited number of latent regression vectors during pretraining learn to perform ridge regression with the Bayes optimal ridge parameter. We instead consider the case where the pretraining task distribution, T Pretrain, contains a limited and finite set of latent regression vectors (see Section 2 for details). To evaluate its ability to learn new tasks, the PT is tested on the ideal task distribution, T True, which is a Gaussian distribution over all latent regression vectors. Studying this setting has two advantages: first, we can directly vary the task diversity in T Pretrain by changing the number of unique latent regression vectors seen during pretraining. Second, we can calculate the optimal estimator that minimizes the pretraining loss-the Bayesian estimator with prior T Pretrain-as well as the optimal estimator for all tasks-the Bayesian estimator with prior T True. This allows us to interpret the behavior of the PT by comparing its predictions to those of the optimal estimators under either task distribution. In our work, we vary the pretraining task diversity and probe the PT's ability to learn fundamentally new tasks in-context: does it behave like the optimal estimator for T Pretrain and perform suboptimally on tasks from T True , or does it align with the optimal estimator for T True which can solve new, unseen tasks?

Contributions. Our contributions are as follows:

We find that a transformer pretrained on data with low task diversity behaves like the Bayesian estimator with prior T Pretrain; it performs optimally on pretraining tasks but cannot learn new tasks in-context. However, as pretraining task diversity increases, the PT deviates from this Bayesian estimator, significantly outperforming it on new tasks, and at a large but still finite number of pretraining tasks, the PT's performance closely matches that of the optimal estimator on T True .

We identify a task diversity threshold for the emergence of ICL. Below this threshold, increasing the pretraining dataset size while keeping task diversity constant biases the PT towards the pretraining task distribution. Conversely, beyond this threshold, increasing the dataset size without increasing its task diversity improves the PT's performance on new, unseen tasks. This suggests that the PT's behavior undergoes a sharp algorithmic phase transition in the limit of many examples per task, aligning with the optimal estimators on T Pretrain before the threshold and on T True after it. We also examine this transition from the perspective of learning dynamics.

We empirically show that increasing the task dimension at fixed SNR increases the task diversity threshold. However, the scaling of the PT's error with dimension is vastly superior to that of the optimal Bayesian estimator for T Pretrain; at a task diversity that is beyond the threshold at all dimensions we consider, the PT remains near-optimal with increasing dimension, whereas the optimal estimator for T Pretrain grows progressively less similar to the optimal estimator for T True .

We show that increasing weight decay significantly decreases the task diversity threshold while increasing number of layers or embedding size increases the task diversity threshold. This elucidates the effect of regularization and model capacity on the emergence of ICL.

Overall these contributions suggest that the emergence of in-context learning in pretrained transformers cannot be fully explained by a theory of Bayesian inference on the pretraining distribution.

2 Problem setup

ICL of linear regression (schematic in Fig. 1). Each ICL task corresponds to a latent D -dimensional regression vector, w ∈ R D . At inference time, the transformer takes as input a sequence of K datatarget pairs, ( x 1 , y 1 , ..., x K , y K ) , which are the in-context examples corresponding to this single task

Transformer

w . For k ∈ { 1 , ..., K } , the data are drawn i.i.d. from a D -dimensional standard normal distribution, x k ∼ N ( 0 , I D ) , and the targets are scalars given by y k = w ⊺ x k + ε k . The ε k 's are noise scalars drawn i.i.d. from a normal distribution with mean 0 and variance σ 2 , ε k ∼ N (0 , σ 2 ) . Let f θ be the PT with parameters θ . Since we use a decoder-only transformer with a causal attention mask, for each k ∈ { 1 , ..., K } , the transformer sees the context S k = ( x 1 , y 1 , ..., x k -1 , y k -1 , x k ) and based on this context, it makes a prediction, f θ ( S k ) , for the target of x k . Thus, in each forward pass, the PT solves K linear regression problems each with the same latent regression vector but an increasing number of in-context examples.

Pretraining. The transformer is pretrained to minimize the next token prediction mean squared error (MSE) on sequences of data and target pairs. The latent regression vector for each sequence is drawn from the pretraining task distribution , T Pretrain. This distribution has limited diversity as it is the uniform distribution over a finite set of M tasks, T Pretrain = U{ w (1) , ..., w ( M ) } . Each task in T Pretrain is drawn i.i.d from a D -dimensional standard normal distribution, w ( i ) ∼ N ( 0 , I D ) , i ∈ 1 , ..., M . By increasing the number of tasks, M , in T Pretrain, we can increase the diversity of the pretraining data. Since the transformer makes a prediction for every data point in the sequence, its loss, L T Pretrain , is just the MSE for each prediction, averaged over the predictions in the sequence:

Evaluation. We evaluate the PT's performance on tasks seen during pretraining by computing L T Pretrain using Eq. (1) but with new samples of data and noise. Since these are new instances of the task with new in-context examples, this evaluation corresponds to the test error of the PT. For a PT to successfully perform ICL of linear regression on new tasks , it must accurately predict the targets from the in-context examples for any task drawn from an ideal task distribution , T True, over all latent regression vectors; in our case T True = N ( 0 , I D ) . We evaluate the PT's performance on new tasks by computing L T True , which follows Eq. (1) but where the tasks are sampled from the ideal task distribution: w ∼ T True in the expectation.

Comparing the PT to optimal estimators. An advantage of studying ICL of linear regression is that we can calculate the ground truth optimal estimators that minimize the loss, L T , in Eq. (1) for both task distributions, T Pretrain and T True. The optimal estimator for the k th prediction, ˆ y T k , that minimizes the k th term in the sum in L T is the Bayesian estimator with T as the prior. This is given by the posterior mean of y k conditioned on the context: ˆ y T k = E T ,ε k [ y k | S k ] , where the expectation is over the task distribution, T , and the noise, ε k (Appendix A.1).

For task distribution T Pretrain = U{ w (1) , ..., w ( M ) } , the discrete minimum mean squared error (dMMSE) estimator is optimal. It is given by ˆ y dMMSE k = ( ˆ w dMMSE k ) ⊺ x k where ˆ w dMMSE 1 =

0.1

0.2

0.3

0.4

Pretrain

MSE/

D

10

5

10

3

10

1

PT, dMMSE

10

4

10

3

10

2

2 0

2 4

2 8

2 12

2 16

2 20

# Pretraining Tasks

0.5

1.0

1.5

True

MSE/

D

dMMSE

PT

2 2

2 4

2 6

2 8

2 10

# Pretraining Tasks

10

5

10

3

10

1

PT, dMMSE

256

512

1024

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

10

4

10

3

10

2

256

512

1024

1 M ∑ M i =1 w ( i ) and for k ∈ { 2 , ..., K } , (Appendix A.2)

Intuitively, ˆ w dMMSE k is just a weighted sum of the pretraining w ( i ) s with weight governed by the likelihood of observing targets { y 1 , ..., y k -1 } conditioned on inputs { x 1 , ..., x k -1 } and the task being w ( i ) . A PT that minimizes the pretraining loss L T Pretrain will behave like this estimator.

For task distribution T True = N ( 0 , I D ) , the Ridge regression estimator with the ridge parameter set to the noise scale σ 2 is optimal: ˆ y Ridge k = ( ˆ w Ridge k ) ⊺ x k , where ˆ w Ridge 1 = 0 and for k = { 2 , ..., K } ,

where X = ( x ⊺ 1 , ..., x ⊺ k -1 ) ∈ R ( k -1) × D and y = ( y 1 , ..., y k -1 ) (Appendix A.3). A PT that performs optimally on new tasks will behave like this estimator. We can compare the behavior of the PT to that of the optimal estimators by computing the mean square difference of the predictions under a given task distribution T . We write this as

3 Experiments and results

Unless specified otherwise, we study linear regression in D = 8 dimensions with up to K = 16 in-context examples and observation noise variance σ 2 = 0 . 25 . We use either a base transformer

2 2

2 4

2 6

2 8

2 10

# Pretraining Tasks

10

5

10

3

10

1

PT, dMMSE

Pretrain

2 2

2 4

2 6

2 8

2 10

# Pretraining Tasks

10

5

10

3

10

1

PT, dMMSE

True

Steps

0.5M

1M

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

10

4

10

3

10

2

Pretrain

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

10

4

10

3

10

2

True

model with the GPT2 architecture [15] with 8 layers, 128-dimensional embeddings, and 2 attention heads or a small model with 4 layers, 64-dimensional embeddings, and 2 attention heads. We train with the Adam optimizer [16] and a one-cycle triangle learning rate schedule [17] with 50% warmup. The base model is trained with batch size 256 for 500K training steps, though these hyperparameters are varied in our experiments. We always sweep over a range of learning rates and choose the largest learning rate at which training is stable. For further details see Appendix B.

To pretrain a randomly initialized transformer on data with task diversity M , we first construct the pretraining task distribution, T Pretrain, as described in Section 2. We then minimize the objective L T Pretrain in Eq. (1) using minibatch stochastic gradient descent. For each sequence in a minibatch, we sample a single task w from T Pretrain, as well as new samples of data, { x i } K i =1 , and noise, { ε i } K i =1 , from their respective continuous distributions, to form a sequence ( x 1 , w ⊺ x 1 + ε 1 , . . . , x K , w ⊺ x K + ε K ) . If we train for N steps at batch size B , the transformer will see a total of NB unique sequences and roughly NB M unique sequences for each latent task in T Pretrain. By increasing either B or N at fixed M , we can increase the total size of the pretraining dataset (or number of sequences per task) while keeping the dataset diversity -the number of unique w s in T Pretrain-fixed.

3.1 Task diversity threshold for the emergence of in-context learning

For Fig. 2, we pretrain our base transformer on datasets with increasing task diversity (on the x-axis) while keeping the total number of sequences seen during pretraining fixed ( B = 256 , N = 500 K ). We evaluate the PTs and both optimal estimators on tasks seen during pretraining drawn from T Pretrain (Fig. 2 top left) and on new tasks drawn from T True (Fig. 2 bottom left) and plot MSE normalized by task dimensionL T /D from Eq. (1)). Since dMMSE is optimal on tasks from T Pretrain (as discussed in Section 2), the green dMMSE markers denote the lowest possible loss the PT could achieve in this setting. In fact, the pretraining objective L T Pretrain explicitly encourages the PT to match dMMSE performance. On the other hand, Ridge is optimal on tasks sampled from T True (Fig. 2 bottom left); the blue markers denote the lowest possible MSE the PT could attain on new tasks.

Low task diversity phase: the PT is Bayesian with respect to the pretraining distribution and cannot solve new tasks. At low pretraining task diversityM up to about 2 6 -the PT's MSE closely tracks that of dMMSE on tasks sampled from T Pretrain (Fig. 2 top left); the PT performs optimally on tasks seen during pretraining. But it significantly underperforms on new tasks sampled from T True, indicated by the gap in MSE between the PT and Ridge (Fig. 2, bottom left). In this regime, it behaves like the Bayesian estimator with prior T Pretrain .

High task diversity phase: the PT is non-Bayesian with respect to the pretraining task distribution and can solve new tasks. At higher task diversities-above 2 6 pretraining tasks-the PT's MSE deviates from dMMSE and approaches Ridge under both T Pretrain and T True. Crucially, the PT starts to significantly outperform dMMSE on unseen tasks sampled from T True (Fig. 2 bottom left) at the expense of not fully minimizing its training objective, L T Pretrain (gap between PT and dMMSE under T Pretrain, Fig. 2 top left). This suggests that, a PT trained on a finite but large number of pretraining tasks can learn fundamentally new tasks in-context and this ability depends on it deviating from the optimal Bayesian estimator on the pretraining task distribution.

0

250K

500K

Step

0

0.02

0.04

0.06

0.08

True

0

1M

2M

Step

Small PT Learning Curves: BS 512, 500K and 2M Steps

#

Pretraining

T asks

2 8

2 9

2 10

2 11

2 12

2 13

2 14

2 5

2 10

2 15

M

1e4

1e5

1e6

t

*

t

*

M

(

=0.47)

t

* = 2

e

6

Scaling of Early Stopping Time

Finite size scaling of training data suggests an algorithmic phase transition as task-diversity increases. The experiments in Fig. 2 (left column) suggest that, when tested on both task distributions T Pretrain and T True, the ICL algorithm implemented by a PT exhibits a smooth crossover in performance from dMMSE to Ridge. We next examine how this transition changes as we increase the number of sequences per task seen over pretraining, at fixed task diversity. One might reasonably expect that, if the transformer sees more sequences per latent task in T Pretrain, both its predictions and performance should become more similar to those of dMMSE, and less similar to those of Ridge, at all values of task diversity. Strikingly, this natural expectation is violated in a manner that facilitates ICL on T True .

At each number of tasks, we increase the number of sequences per task by increasing batch size from 256 to 512 to 1024, while leaving the number of training steps fixed at 500K. We observe that ∆ T Pretrain PT, dMMSE , which quantifies how different the PT and dMMSE estimator's predictions are when testing on tasks drawn from T Pretrain, does in fact decrease for M ≤ 2 10 (Fig. 2 top center) as we train on more sequences per task. Moreover, for each M ∈ { 2 10 , ..., 2 14 } the PT's predictions also become less similar to those of Ridge, both on tasks from T Pretrain (Fig. 2, top right) and T True (Fig. 2, bottom right). Crucially , this movement in behavior of the PT towards dMMSE and away from Ridge, at least on tasks drawn from T Pretrain, holds only up to a threshold number of tasks between 2 14 and 2 15 . Beyond this threshold, pretraining on more sequences per task at a fixed task diversity actually makes the PT more like Ridge, in that both ∆ T Pretrain PT,Ridge and ∆ T True PT,Ridge decrease (Fig. 2, right top and right bottom respectively). This means that, beyond a task diversity threshold, the PT can not only optimally solve new tasks from T True by matching Ridge performance, but also the PT gets better at doing so if trained on more sequences per task, despite the limited set of tasks experienced in pretraining. Thus, in contrast to the natural expectation stated above, more sequences per task does not promote overspecialization of the PT to the T Pretrain at task diversities beyond the threshold.

Finally, the motion of the ICL algorithm implemented by PT towards (away) from Ridge above (below) a task diversity threshold (Fig. 2, right top and bottom) indicates that as one increases the number of sequences per task at fixed task diversity, the smooth cross over in performance of the PT between dMMSE and Ridge, shown in Fig. 2, left top and bottom, will become sharper and sharper in task diversity, ultimately exhibiting a sharp phase transition in the limit of infinite number of sequences per task. Remarkably, this phase transition in the ICL algorithm implemented by the PT appears at a moderate task diversity threshold below 2 15 pretraining tasks; even though dMMSE significantly underperforms relative to Ridge on T True at this task diversity, the PT nevertheless remains unimpaired by this limited task diversity and can optimally solve new tasks.

Increased training time at fixed batch size further supports an algorithmic phase transition. To confirm the above results, we also increase the number of sequences per task, at each task diversity, by increasing the number of training steps N from 500K to 1M while keeping batch size fixed at 256. We observe that doubling N (change from pale blue to red in Fig. 3) and doubling B (change from pale

0.0

0.2

0.4

0.6

0.8

1.0

Interpolation

0.2

0.3

0.4

0.5

0.6

MSE/

D

2 5  Pretraining Tasks

0.0

0.2

0.4

0.6

0.8

1.0

Interpolation

2 10  Pretraining Tasks

0.0

0.2

0.4

0.6

0.8

1.0

Interpolation

2 15  Pretraining Tasks

dMMSE

PT

blue to red in Fig. 2) have very similar effects on ∆ T PT,dMMSE and ∆ T PT,Ridge , for both T = T True and T = T Pretrain. More importantly, the task diversity threshold, which we determined as the cross-over point in ∆ T True PT,Ridge between batch sizes 256, 512, and 1024 at 500K training steps (Fig. 2 bottom right) happens at the same number of tasks as the crossover point between 500K and 1M steps at batch size 256 (Fig. 3, right). Given that our two approaches for training the baseline transformer on more data yield the same task diversity threshold, and that doubling batch size leads to significantly faster training times than doubling number of steps, from here onward we consider the task diversity threshold to be cross-over point in ∆ T True PT,Ridge between batch sizes 256 and 512 when training for 500K steps. See Appendix D for more ablations of batch size and training steps that provide further evidence for how the number of sequences seen by the transformer is the key factor determining the similarity of its predictions to those of dMMSE and Ridge at each number of tasks.

Learning dynamics and a break in the scaling of early stopping time further supports an algorithmic phase transition. To probe if the observed transition is merely an effect of under-fitting, we study the learning dynamics of small PTs in the very large number of steps regime. First, in Appendix E, we verify that the small PT also demonstrates an algorithmic phase transition but at lower task diversity threshold between 2 11 and 2 12 pretraining tasks. In Fig. 4 left, we visualize the learning curves ( ∆ T True PT,Ridge vs training steps) of PTs trained for 500K steps at batch size 512 with pretraining task diversities, M , below and above the task diversity threshold, M ∗ . For M < M ∗ , ∆ T True PT,Ridge decreases early in training until it reaches a minimum at time step t ∗ , and then increases as the PT approaches dMMSE. We define t ∗ as the early stopping time for Ridge. For M > M ∗ , ∆ T True PT,Ridge decreases throughout training. To evaluate if, in the latter case, models are undertrained and t ∗ is larger than the total training time, we extend training to 2M steps at batch size 512 ( 4 × the training time, see Appendix B). Fig. 4 center, shows these learning curves along with that of the model trained with infinite task diversity; even in this long training regime, the task diversity threshold does not change. For both short and long training durations, models trained with the same M have similar qualitative behavior (whether distance to Ridge decreases then increases or monotonically decreases). Additionally, learning curves of the models with M > M ∗ are very similar to the learning curves for models trained on infinite pretraining task diversities and they achieve similar final accuracy (dahed lines vs markers in Fig. 10), suggesting that these models are approaching the Ridge solution.

In Fig. 4 right, we study how t ∗ , scales with M . For most M < M ∗ , t ∗ obeys a simple scaling behavior t ∗ ∝ M α , α ≈ 0 . 47 . However, for M > 2 10 , the distance to Ridge decreases monotonically through training and t ∗ = 2M steps. Despite the caveat that our experiments are necessarily in the large but finite training step regime with a decayed learning rate schedule, this stark break in the scaling behavior of t ∗ near the task diversity threshold suggests that the observed transition is not just caused by under-fitting but an underlying difference in the learning dynamics.

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

10

4

10

3

10

2

10

1

True

8 Dimensions

256

512

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

10

4

10

3

10

2

10

1

16 Dimensions

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

10

4

10

3

10

2

10

1

32 Dimensions

8

16

24

32

Task Dimension

0.0

0.1

0.2

0.3

0.4

0.5

M

=2 20

Model

dMMSE

PT

The transition along interpolating paths. To obtain an additional description of the algorithmic transition in the PT from dMMSE to Ridge, we compute the ICL performance of the PT, and compare it to both dMMSE and Ridge, on a one parameter family of new tasks w α that interpolate between pairs of seen tasks w i and w j in the support of T Pretrain. The interpolation path is given by

Here we fix the norm of the interpolated vector w α to the average of the two endpoint norms to avoid ∥ w α ∥ taking on very small values for α ∼ 1 2 . Fig. 5 shows the results of this analysis for 2 5 (left, low task diversity regime), 2 10 (center, below task diversity threshold), and 2 15 (right, just above the task diversity threshold) tasks. At each value of α , MSE is averaged over a large number of task pairs ( w i , w j ) . Examination of the average performance at the center of many interpolation paths, corresponding to fundamentally new tasks far from tasks seen during pretraining, clearly reveals a transition in PT performance from dMMSE to Ridge, where new tasks can only be optimally learned above, but not below, the task diversity threshold. In contrast, unlike the PT, dMMSE cannot solve new tasks at any task diversity in the range considered.

The PT outperforms a smoothed dMMSE model. We have seen that at an intermediate task diversity the PT significantly outperforms dMMSE on new tasks in T True. It is clear why dMMSE performs poorly on new tasks in T True at low task diversity: its prior over tasks concentrates on M unique tasks in T Pretrain, while the prior over tasks in T True is Gaussian. A natural conjecture is that the PT cannot memorize all M tasks in T Pretrain for large enough M . Therefore we also compare PT performance to a smoothed dMMSE estimator in which the discrete point prior over M tasks seen in pretraining is replaced with a mixture of M isotropic Gaussians with the same centers but with variance chosen to optimize performance on T True (see Appendix G for details). This smoothed dMMSEoutperforms dMMSE as it has a prior over tasks closer to the Gaussian T True. But remarkably, the PT still outperforms the smoothed dMMSE even with optimal smoothing (Fig. 12). This indicates the PT, even at moderate task diversity, implements a more sophisticated algorithm than a simple smoothed dMMSE arising from the PT's inability to resolve the M pretraining tasks to high precision.

3.2 The PT exhibits superior scaling of task diversity threshold with dimension than dMMSE.

We next explore the dependence of the task diversity threshold on the regression problem dimension D . We vary D ∈ { 8 , 16 , 32 } while simultaneously scaling up maximal context length as K = 2 D , and increasing observation noise σ 2 to match the SNR to that of D = 8 and σ 2 = 0 . 25 . We also train a larger model with 12 layers, 256-dimensional embeddings, and 4 attention heads that is sufficiently expressive to match Ridge performance at D = 32 . Fig. 6, first 3 panels reveal that the task diversity threshold of the PT increases moderately (approximately linearly) with task dimension (i.e. roughly 2 14 , 2 15 , and 2 16 at D = 8 , 16 , and 32 respectively). This linear scaling is remarkable considering the volume of all possible tasks scales exponentially with dimension due to the concentration of the Gaussian T True to a sphere for large D . Thus we expect dMMSE performance to scale much more poorly with dimension D since the finite number of tasks in T Pretrain would need to cover a substantial portion of the sphere for dMMSE to approach Ridge. To test this hypothesis, for M = 2 20 , which is

0

0.01

0.1

1

Weight Decay

2 10

2 11

2 12

2 13

2 14

2 15

Task Diversity Threshold

Base

Small

Ablate Weight Decay

64

128

256

512 1024

Embedding Dimension

2 11

2 12

2 13

2 14

2 15

Task Diversity Threshold

Ablate Embedding Dimension

4

6

8

10

# Layers

2 11

2 12

2 13

2 14

Task Diversity Threshold

Ablate # Layers

the largest task diversity we consider, we explore how the similarity of PT and dMMSE predictions to Ridge on new tasks scales with D (Fig. 6, right panel). We see that ∆ T True dMMSE,Ridge grows significantly as we increase D , while remarkably ∆ T True PT,Ridge is largely dimension independent. Overall this indicates that the scaling of PT error with dimension is vastly superior than that of dMMSE; PT remains near optimal and close to Ridge at all D for M = 2 20 , while dMMSE departs from Ridge as D increases.

3.3 Effect of Regularization and model capacity on the task diversity threshold.

We study the dependence of the task diversity threshold on various hyperparameters. First, adding explicit regularization in the form of weight decay (see Appendix B for details), and increasing its value over three orders of magnitude, consistently lowers the threshold task diversity (Fig. 7, left). Note however, the lower task diversity threshold also comes with worse performance (Figure 13, top). This suggests that various forms of implicit regularization could help drive the algorithmic transition in the PT without weight decay. We also explore the effect of model capacity on the task diversity threshold by either increasing the embedding dimension of both small and base PTs or increasing the depth of small PTs. Fig. 7 center shows that increasing embedding dimension over a reasonable range does not affect the task diversity threshold of base PT. However, for small PT, increasing either the embedding dimension (Fig. 7 center) or depth (Fig. 7 right) increases the task diversity threshold. Base PT has a much larger capacity then small PT and also has a larger threshold; we hypothesize that small PT is still in a regime where the threshold is sensitive to capacity while base PT is not. Together, these results suggest that model capacity plays an important role in the emergence of in-context learning: increasing capacity (up to a point) leads to an increase in the task-diversity threshold.

4 Related work

The Bayesian framework for ICL introduced by Xie et al. [3], which motivates our work, hypothesizes that PTs "locate" concepts learned during pretraining to solve ICL tasks. A series of empirical work in language models [18-20] use this framework to select better in-context examples while Min et al. [21] use it to study the robustness of latent task inference. Our work builds on this framework in the linear regression setting and validates it at low task diversities. However, we find a regime-large but finite number of pretraining tasks-in which the ability to learn new tasks in-context is an emergent phenomenon that cannot be fully explained by Bayesian inference.

Prior work [12-14] has also shown that transformers can do linear regression in-context. However, they pretrain with unlimited task diversity, sampling a completely new regression vector for each sequence. In contrast, our work considers pretraining datasets with limited task diversity where ICL on new tasks emerges even though the pretraining loss does not explicitly encode it. Another line of work hypothesizes that ICL performs gradient descent in the activations of the forward pass, providing explicit constructions for the weights of the PT to implement this for linear regression [13, 14] or exploring this hypothesis in language models [22]. However more experiments are required to test the hypothesis that trained transformers actually match proposed constructions. Instead of studying the explicit mechanism by which in-context learning is implemented, our work focuses on the impact of the pretraining task diversity. Similar questions pertaining to the role of task diversification have been explored in the meta-learning literature [23-25].

Kirsch et al. [26] also show the emergence of in-context learning with pretraining task diversity on a toy classification task. By studying this question in the controlled setting of linear regression, we can compare to the optimal estimators on T Pretrain and T True. This allows us to establish that ICL at finite task diversity emerges because the PT departs from the optimal estimator on the pretraining task distribution, and is not just a consequence of the pretraining task distribution becoming similar to the ideal task distribution. Among other important perspectives on ICL, Chan et al. [27] identify, in a toy setting, several properties of the training distribution-burstiness and occurrence of rare classes-that are necessary for the emergence of ICL. Wei et al. [28] study how ICL in large language models is affected by semantic priors and input-label mappings, focusing on differences across model scale. Olsson et al. [29] study inductions heads-circuits responsible for completing patterns by copying tokens-as a mechanism for implementing ICL.

5 Discussion

Overall, we have extensively explored the impact of pretraining task diversity on the emergence of incontext learning of fundamentally new tasks not seen during pretraining. We found several surprises by working in the controlled setting of linear regression, where we could compare the performance of the PT to Bayesian estimators that are optimal, either for the limited diversity pretraining task distribution T Pretrain (i.e. dMMSE), or for the diverse ideal task distribution T True (i.e. Ridge). These comparisons reveal an algorithmic phase transition in the PT from the former to the latter at an intermediate task diversity threshold; beyond this threshold, the PT solves fundamentally new tasks not seen during pretraining. Strikingly, this task diversity threshold scales moderately with task dimension, over the range of dimensions considered, despite the exponential growth in the volume of all possible tasks with dimension. Indeed this PT scaling vastly outperforms that of dMMSE. Overall, these results indicate that ICL of new tasks by PTs is an emergent phenomenon that cannot be explained by Bayesian inference on limited diversity pretraining task distributions. Moreover, our experiments suggest some form of implicit regularization in PTs allows them to break free of the pretraining task distribution to solve new tasks, given a moderate pretraining task diversity.

Remarkably, beyond the task diversity threshold, PTs learn the optimal estimator for the underlying generative model for pretraining tasks; this is the case for both Gaussian and Laplace priors over tasks (see Figure 14 for experiments with Laplace prior). This is true even though solutions with lower training loss exist; indeed when trained on more data at fixed diversity, PTs behave more like Ridge at the expense of higher training loss. Our experiments in Fig. 4 suggest that this algorithmic transition is due an underlying change in learning dynamics. We explore this hypothesis by probing the linear mode connectivity of the loss landscape [30, 31]. In Fig. 11 we find that PTs trained with large M inhabit the same loss basin as PTs trained with M = ∞ : the training loss barrier between PTs trained with M ≳ 2 13 and PTs with M = ∞ is similar to two PTs trained with M = ∞ . In contrast, there are large loss barriers between PTs trained with M < 2 13 and M = ∞ . Additionally, PTs trained with M = ∞ are closer in weight space to PTs trained with large M than those trained with small M (see Appendix F). Overall, these experiments provide further evidence that PTs trained with task diversities beyond the threshold find solutions similar to the optimal model for T True; we leave further exploration of these loss landscapes to future work.

An intriguing question is how these observations carry over to language. A key mystery about the efficacy of ICL in language tasks lies in how different the tasks learned in-context are from the pretraining distribution of large language corpora. It is also less clear how to categorize the contents of such corpora according to tasks and measure their resulting task diversity. Regardless, our observation in linear regression that a moderate threshold in pretraining task diversity can enable PTs to solve new tasks may imply that many language tasks that are quite different from the statistics of large language corpora can still nevertheless be solved in-context.

Our results also suggest that the scale of data alone does not lead to good ICL performance. In fact, below the task diversity threshold, increasing the size of the pretraining dataset without increasing task diversity hurts ICL performance. It is necessary to increase both the diversity and size of the dataset for ICL to emerge. Thus to improve ICL in language settings, our work motivates future studies into uncovering the relevant notion of tasks in language modeling and approaches to increase task diversity in language corpora. More generally, our empirical analysis of the impact of pretraining task diversity on ICL motivates further theoretical studies. Such studies will be key to understanding the mystery of why simple next token prediction during pretraining can lead to in-context learning of so many apparently different tasks.

References

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ 1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf .

Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant. Risks from learned optimization in advanced machine learning systems, 2021.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 , 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 21(1), jan 2020. ISSN 1532-4435.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=gEZrGCozdqR .

Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szczechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, Debajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang, Han Wang, Matteo Manica, Sheng Shen, Zheng Xin Yong, Harshit Pandey, Rachel Bawden, Thomas Wang, Trishala Neeraj, Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Teehan, Teven Le Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=9Vrb9D0WI4 .

Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap, Eshaan Pathak, Giannis Karamanolakis, Haizhi Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby Kuznia, Krima Doshi, Kuntal Kumar Pal, Maitreya Patel, Mehrad Moradshahi, Mihir Parmar, Mirali Purohit, Neeraj Varshney, Phani Rohitha Kaza, Pulkit Verma, Ravsehaj Singh Puri, Rushang Karia, Savan Doshi, Shailaja Keyur Sampat, Siddhartha Mishra, Sujan Reddy A, Sumanta Patro, Tanay Dixit, and Xudong Shen. Super-NaturalInstructions: Generalization via declarative instructions on 1600+ NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 5085-5109, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https: //aclanthology.org/2022.emnlp-main.340 .

Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetaICL: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 2791-2809, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/ 2022.naacl-main.201. URL https://aclanthology.org/2022.naacl-main.201 .

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 27730-27744. Curran Associates, Inc., 2022. URL https://proceedings.neurips.cc/paper_files/paper/ 2022/file/b1efde53be364a73914f58805a001731-Paper-Conference.pdf .

Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Yi Tay, Denny Zhou, Quoc V Le, Barret Zoph, Jason Wei, et al. The flan collection: Designing data and methods for effective instruction tuning. arXiv preprint arXiv:2301.13688 , 2023.

Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=flNZJ2eOet .

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/ forum?id=0g0X4H8yN4I .

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2022.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR , abs/1412.6980, 2014.

Leslie N. Smith and Nicholay Topin. Super-convergence: very fast training of neural networks using large learning rates. In Defense + Commercial Sensing , 2017.

Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for incontext learning. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 2655-2671, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/ 2022.naacl-main.191. URL https://aclanthology.org/2022.naacl-main.191 .

Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for GPT-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures , pages 100-114, Dublin, Ireland and Online, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.deelio-1.10. URL https: //aclanthology.org/2022.deelio-1.10 .

Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning, 2023.

Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 11048-11064, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https://aclanthology.org/2022.emnlp-main.759 .

Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models implicitly perform gradient descent as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models , 2023. URL https://openreview.net/forum?id=fzbHRjAd8U .

Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Metalearning without memorization. 2020.

Ramnath Kumar, Tristan Deleu, and Yoshua Bengio. The effect of diversity in meta-learning. 2022.

Nilesh Tripuraneni, Chi Jin, and Michael I. Jordan. Provable meta-learning of linear representations, 2022.

Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose incontext learning by meta-learning transformers, 2022.

Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya Singh, Pierre H. Richemond, Jay McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers, 2022.

Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently, 2023.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread , 2022. https://transformer-circuits.pub/2022/in-context-learning-and-inductionheads/index.html.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis, 2020.

Mansheej Paul, Feng Chen, Brett W. Larsen, Jonathan Frankle, Surya Ganguli, and Gintare Karolina Dziugaite. Unmasking the lottery ticket hypothesis: What's encoded in a winning ticket's mask? In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=xSsW2Am-ukZ .

A Bayesian estimators

A.1 Bayesian MMSE estimator

The optimal estimator for the k th prediction, ˆ y T k , that minimizes the k th term in the sum in L T (in Eq. (1)) is the Bayesian estimator with T as the prior. This is given by the posterior mean of y k conditioned on the context: ˆ y T k = E T ,ε k [ y k | S k ] , where the expectation is over the task distribution, T , and the noise, ε k . To derive this, notice that, using the law of total expectation, the k th term in the loss L T is:

which implies that the optimal estimator is ˆ y T k = E T ,ε k [ y k | S k ] . Adopting the notation X = ( x ⊺ 1 , ..., x ⊺ k -1 ) ∈ R ( k -1) × D and y = ( y 1 , ..., y k -1 ) , we can express the expectation explicitly:

A.2 dMMSE estimator

For task distribution T Pretrain = U{ w (1) , ..., w ( M ) } , we can directly get the discrete minimum mean squared error (dMMSE) estimator by plugging the uniform discrete distribution into Eq. 6.

A.3 Ridge estimator

For task distribution T True = N ( 0 , I D ) , we can directly get the Ridge regression estimator from Eq. 6:

B Experimental details

For most experiments, we study linear regression in D = 8 dimensions with up to K = 16 in-context examples and observation noise variance σ 2 = 0 . 25 . Our base model is a transformer with the GPT2 architecture [15] with 8 layers, 128-dimensional embeddings, and 2 attention heads. We train with where

the Adam optimizer [16] and a one-cycle triangle learning rate schedule [17] with 50% warmup. The base model is trained with batch size 256 for 500K training steps, though these hyperparameters are varied in our experiments. Specifically, for the experiments in Fig. 3 we sweep over 500K and 1M training steps to find the task diversity threshold as a function of training steps. For all other experiments, we sweep batch size 256 and 512 and for the experiments in Fig. 2, we do an additional batch size of 1024.

For the experiments in Fig. 6, we use a larger model with 12 layers, 256-dimensional embeddings, and 4 attention heads. This ensures that the model can learn to perform ridge regression in the D = 32 setting. We also train the models on D = 8 , 16 , 24 , 32 with up to K = 2 D in-context examples in each case. The observation noise variance is scaled at each dimension to keep the signal to noise ratio constant. So σ 2 = 0 . 5 , 0 . 707 , 0 . 866 , 1 for D = 8 , 16 , 24 , 32 respectively.

We always sweep over learning rates in {0.0001, 0.0003, 0.001} and choose the largest learning rate at which training is stable. For almost all experiments, the learning rate is 0.001, the main exception being the experiments in Fig. 6 where the learning rate is 0.0001.

For the 2M training step experiments in Fig. 4 and Fig. 10, the learning rate increases linearly up to 0.001 at 250K steps (i.e. using the same warmup as in the 500K step experiments) and then decreases linearly to 0 at 2M steps, consistent with our choice of decayed learning rate schedules. The experiments in Appendix F are performed at batch size 512 and following the 500K step one-cycle triangle learning rate schedule described above.

The implementation for this paper was done in JAX and all experiments were run on TPU v2-8 and v3-8 provided as part of the Google TPU Research Cloud program. Each training run takes approximately 4 v2-8 TPU hours.

C Support Figure 2

Fig. 8 provides an additional visualization for Fig. 2 upper middle but on a linear scale on the y-axis. In this figure it is clear that for all task diversities below the threshold-less that 2 14 -increasing the batch size aligns the PT with dMMSE, the optimal estimator for T Pretrain. Thus below the task diversity threshold, the optimal estimator does indeed approach the Bayesian estimator with prior T Pretrain when trained on more data and is thus unable to learn new tasks as discussed in Section 3.1.

2 2

2 4

2 6

2 8

2 10

2 12

2 14

# Pretraining Tasks

0.00

0.02

0.04

0.06

0.08

Pretrain

PT, dMMSE

256

512

1024

D Dependence on number of sequences per task

In Section 3.1 we explore how the crossover of the PT from behaving like dMMSE to behaving like Ridge as we increase pretraining task diversity depends on the number of sequences per task seen during pretraining. In Fig. 2 middle and right column, we probe this by increasing the batch size to increase the number of sequences per task seen during pretraining at each level of task diversity. To ensure that that the phase transition and task diversity threshold we find is indeed a consequence of increasing the number of sequences per task and not merely just an effect of increasing batch size, we keep the batch size constant and increase the number of sequences per task in Fig. 3.

In Fig. 9 we provide additional evidence that the number of sequences per task and not batch size nor number of training steps is the primary factor that drives the behavior of the final model. Specifically, we show that the final behavior of the model is invariant to batch size or number of training steps if the number of training sequences per task is kept constant, atleast within a reasonable range of batch size and number of training steps. To do so, at each level of pretraining task diversity, we compare our base model (batch size = 256, 500K training steps, red cicles) to a model trained with batch size 512 (light blue squares) and batch size 128 (purple triangles). For either variant of the batch size, we adjust the number of training steps to keep the total number of pretraining sequences per task constant: 250K training steps for batch size 512 and 1M training steps for batch size 128. In all three cases, the behavior of the PT is identical as measured by ∆ T PT,dMMSE and ∆ T PT,Ridge for both T Pretrain and T True across all numbers of pretraining tasks.

2 2

2 4

2 6

2 8

2 10

# Pretraining Tasks

10

5

10

3

10

1

PTM,dMMSE

Pretrain

2 2

2 4

2 6

2 8

2 10

# Pretraining Tasks

10

5

10

3

10

1

PTM,dMMSE

Latent

512, 250K

256, 500K

128, 1M

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

10

4

10

3

10

2

Pretrain

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

10

4

10

3

10

2

Latent

E Small PT

We run our task diversity threshold finding experiment in a small PT with 64 dimensional embeddings, 4 layers, and 2 attention heads. The existence of a task diversity threshold reproduces in a Small PT: the PTs transition from becoming less like Ridge to becoming more like Ridge when trained on more data (see Fig. 10). We also note that the Small PT with a lower capacity has a lower threshold ( ∼ 2 11 . 5 compared to ∼ 2 14 . 5 for the Base PT in Fig. 2) but also has lower overall performance (compare y-axis range to Fig. 2 lower right).

2 10

2 12

2 14

2 16

2 18

M

(# Pretraining Tasks)

10

3

10

2

True

Small PT

256, 500K

512, 500K

512, 2M

M

=

barrier

1.5

M=

180

M=∞

Mfinite

distance

M finite

1.0

160

loss

140

Tpretrain

0.5

5

TTTT

120

天天天天天天工T

0.0

100

10

15

10

15

log2 (M)

log2 (M)

F Examining the loss landscape along paths interpolating between finite and infinite task models

For each finite number of tasks, M , as well as M = ∞ (where each training example samples a new task from the Gaussian distribution), we train four small PTs, each with a different random seed for tasks . This means that all randomness (including model initialization, data, and observation noise) is shared across the four experiments, except for that governing the sampling of w 's used for training. We then compare the finite M models to the M = ∞ models, both in terms of the size of the training loss barrier along a path interpolating between the models, as well as the distance in weight space between the models.

We compute the training loss barrier between a finite M model and an infinite task model under the objective L T Pretrain defined by the specific set of M tasks the finite task model was trained on. We follow the approach outlined in [31] where the barrier between two models is the loss at a model whose weights are the average of the two models, minus the average of the losses of the two models. Specifically, given two weight configurations w 1 and w 2 and a loss function L , the loss barrier is calculated as

The distance in weight space, in turn, is simply the L 2 norm of the difference in model weights. For each M we average barriers and distances across all finite-infinite task pairs of models. We treat the analogous quantities computed between pairs of infinite task models as a baseline. See Fig. 11 and the Discussion section in the main text.

G Smoothed dMMSE

2 12

2 14

2 16

2 18

2 20

# Pretraining Tasks

0.345

0.346

0.347

MSE/

D

True

Smooth dMMSE

PT

Here we consider the setting in which the discrete point prior over M tasks seen in pretraining is replaced with a mixture of M isotropic Gaussians with variance ϵ 2 . We call the estimator that does posterior predictive inference under this prior the Smoothed dMMSE estimator (sMMSE). In the limit of ϵ → 0 , the sMMSE estimator becomes the dMMSE estimator. On the other hand, if ϵ →∞ , the prior will look like a Gaussian distribution but will have a variance much larger than the variance of T True. In both edge cases, it won't perform well on T True, and so there is an optimal ϵ that achieves the best performance on T True. In practice, we perform a search by simulation to determine the optimal ϵ . We observe in Fig. 12 that for 2 13 tasks onward, the PT outperforms the optimally smoothed sMMSE on T True .

To derive the sMMSE estimator, we want to calculate ˆ w k = E [ w | S k ] and apply Eq. (6).

We define, so that,

Plugging back to Eq. (6),

where β = softmax ( ˜ β ) and ˜ β i ≡ 1 2 J ⊺ i A -1 J i -∥ w ( i ) ∥ 2 2 ϵ 2 .

H Supporting figures

2 10

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

10

1

10

2

10

3

10

4

True

Weight Decay: 0.0

256

512

2 10

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

Weight Decay: 0.01

2 10

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

Weight Decay: 0.1

2 10

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

Weight Decay: 1.0

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

10

4

10

3

10

2

10

1

True

128 Embedding

256

512

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

256 Embedding

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

512 Embedding

2 12

2 14

2 16

2 18

2 20

# Pretraining tasks

1024 Embedding

Laplace prior

BatchSize

TTrue

0.345

256

512

on

M=∞

0.340

MSE/D

0.335

212

214

216

218

# Pretraining tasks

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper empirically studies whether the ability of transformers to perform in-context learning can be explained by Bayesian inference in a setting wherein tasks are linear regression problems. The pre-training task distribution is the uniform distribution over M unique tasks (ground-truth regressors), and the population distribution of tasks is the standard Gaussian distribution over ground-truth regressors in $R^D$. The experiments shed light on a number of surprising phenomena, all centered around the observation that the estimator found by pre-training a transformer with in-context learning (PT) behaves similarly to the optimal Bayesian estimator with prior given by the pre-training task distribution when $M$ is below a threshold, but when $M$ exceeds this threshold, PT behaves more similarly to the optimal Bayesian estimator with prior being the population distribution of tasks. Moreover, experiments show that this threshold appears to grow only linearly with data dimension $D$, and can be reduced via weight decay. Additionally, when $M$ exceeds the threshold, using more samples from the pre-training tasks during pre-training brings PT closer to the optimal Bayesian estimator with population prior, and increasing dimension $D$ does not significantly affect the behavior of PT, while it severely degrades the performance on the population task distribution of the optimal Bayesian estimator with prior being the pre-training task distribution.\n\nSTRENGTHS:\n1. The paper addresses the highly relevant problem of understanding the ability of large transformers to perform in-context learning, of which little is known.\n\n2. The paper is very well written. The motivation is strong and the formulation is clear. The experiments are well-explained. I did not notice any typos.\n\n3. The experiments are rigorous.\n\n4. The setting is simple and easy to understand, and simultaneously the experiments reveal a number of surprising phenomenon relevant to understanding in-context learning. Thus, this paper should open doors to a number of important follow-up theoretical analyses.\n\nWEAKNESSES:\n1. A downside to the simple formulation -- it is not clear whether the observations here generalize to the large language models and textual data used in practice, which the authors acknowledge. Nevertheless, we must understand this simple setting before generalizing to more complex scenarios.\n\n2. The paper does not provide any theoretical analysis even though the setting is ripe with open theoretical questions and possibly tractable answers. However, the experimental contributions alone are novel and significant enough to merit acceptance in my opinion.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the generalization abilities of a next token prediction pre-training with Transformers architecture to solve new unseen tasks. It performs a series of controlled experiments on simple multivariate regression tasks. The authors show that, unlike suggested in previous works, the solution learned by the transformer architecture deviates from that of the Bayesian solution with finite-support prior over tasks (dMMSE) and converges to the Ridge solution in the finite-set regime past a certain threshold. They empirically characterise when this transition occurs and study how different factors influence it.\n\nSTRENGTHS:\n1. The paper is well-written and introduces all the necessary background needed for understanding.\n2. The paper studies the generalization of transformers pre-trained with next token prediction to solving novel tasks in-context in a few-shot manner which is an important direction to understand how models such as GPT generalize to solve language tasks seemingly different from what they were trained on.\n3. The empirical study is well controlled and supports the conclusions made.\n\nWEAKNESSES:\n1. While described by a new term, “in-context learning” (ICL), the setting studied in the paper is practically identical to that of few-shot and meta-learning. Therefore, it is important to position w.r.t. this line of work, including the studies on the role diversification of pre-training tasks, e.g., [1-3]. See also [4] for theoretical foundations on learning to learn.\n2. The problem of underfitting is not studied sufficiently well. It is mentioned in the paper that, in the main experiment (Fig. 1), the PT model underfits as the error on the training tasks is higher than that of the optimal solution. This would be an important aspect to clarify as one would expect that a model with sufficient capacity should converge to the optimal solution on the training dataset, i.e., dMMSE.\n\n[1] M. Yin, G. Tucker, M. Zhou, S. Levine, and C. Finn, “Meta-Learning without Memorization.” arXiv, Apr. 27, 2020. doi: 10.48550/arXiv.1912.03820.\n\n[2] R. Kumar, T. Deleu, and Y. Bengio, “The Effect of Diversity in Meta-Learning.” arXiv, Nov. 24, 2022. doi: 10.48550/arXiv.2201.11775.\n\n[3] K. Hsu, S. Levine, and C. Finn, “Unsupervised Learning via Meta-Learning.” arXiv, Mar. 21, 2019. doi: 10.48550/arXiv.1810.02334.\n\n[4] J. Baxter, “Theoretical Models of Learning to Learn,” in Learning to Learn, S. Thrun and L. Pratt, Eds., Boston, MA: Springer US, 1998, pp. 71–94. doi: 10.1007/978-1-4615-5529-2_4.\n\nQUESTIONS:\n1. (See point in weaknesses for more context). Is the underfitting seen in Fig.1 top left due to the evaluation on the new data points (for the same tasks) or insufficient model capacity? If latter, did authors try to increase capacity with parameters other than the dimensionality of the embedding in Sec 3.3, e.g., the depth/number of attention heads? It would be interesting to see whether there is an implicit bias in the transformer architecture (trained with SGD) regardless of its capacity.\n    - Note that this observation also contradictions to the statement made in L124: \"A PT that minimizes the pretraining loss LTPretrain will behave like this estimator.\"\n2. What are “fundamentally” new tasks? This term is mentioned multiple times, but it is not defined what makes some tasks “fundamentally” new. Moreover, the new test tasks are sampled from the same distribution as the training ones ($w_i \\sim N(0, I)$). See the next question for a broader discussion.\n3. What makes the observed behavior (the transition from dMMSE -> Ridge) emergent and different from a superior iid generalization of one algorithm (next token prediction with Transformers) over the other (Bayesian dMMSE)?\n    - In standard supervised learning, iid generalization occurs from training to test examples sampled from the same distribution. In this case, one learning algorithm, e.g., ResNet trained with SGD, can exhibit better iid generalization than another, e.g., AlexNet, i.e., converges to the true solution with fewer training examples. \n    - Similarly, the paper shows that one (meta-)learning algorithm (next token prediction with Transformers) is superior to the Bayesian dMMSE algorithm.\n    - More specifically, the setting of ICL can be described in terms of standard supervised learning (perhaps, with varying dimensionality of input). Indeed, the input can be defined as $x = S_k$ and the output as $y_k$, then training examples are $(S_k, y_k)$.\n    - Are these analogies (fairly) the same, or is there a fundamental difference?\n\nLIMITATIONS:\n1. Similar to the raised question about the similarity to a standard iid generalization setting, the experimental setup of this paper does not consider the distribution shift occurring between train and test tasks, an important aspect of generalization (most probably) occurring between LM pre-training and natural language tasks.\n    -  Sec 5 of the paper mentions that the diversity of the pre-training dataset of tasks is important. However, this is only supported when more tasks are sampled from the same distribution as test tasks ($w \\sim N(0, I)$ in both train and test). Whereas in language, there is (most probably) a distribution shift between them.\n2. The paper limits the study to transformer architectures. It would be interesting to see how other architectures used in few-shot and meta-learning behaves in the same settings, e.g., [1]\n\n[1] M. Garnelo et al., “Conditional Neural Processes.” arXiv, Jul. 04, 2018. doi: 10.48550/arXiv.1807.01613.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates in-context learning in the restricted setting of linear regression. Diverging from recent work, they focus on studying the  effect of task diversity i.e. the number of different linear teachers used during pre-training of Transformer models. \nSurprisingly, they find that although fitting the training data optimally, the Transformer outperforms a Bayesian estimator in an intermediate stage of diversity coming closer to the optimal solution of ridge regression. This implies an implicit regularization in the pre-training of the Transformer favoring the better generalizing solution.\n\nSTRENGTHS:\nThe paper studies a highly relevant problem, that of in-context learning, in a restricted setting of linear regression. I find the paper very well written, easy to follow and insightful. Well done!\n\nWEAKNESSES:\nThe major concern about this paper is that it provides only minor novel insights in my opinion. We know that Transformers when pre-trained on enough data implement functions that resemble gradient descent and act as the ridge regression solution. Now given a fixed amount of \nteacher weights, if the Transformer is large enough simply learn and remember all needed weights (\"features\") and implement dMMSE. \nNow, when the capacity is not sufficient anymore it falls back to ridge regression which happens to generalize better. \nIn my opinion, this just confirms classic statistical learning theory, here with the nice twist of Transformers implementing algorithms on both extremes. Please correct me if I am wrong here. \nAlthough I find the paper still interesting, for me to give a higher score various questions and experiments are needed.   \n\nI have a couple of open ended questions on which I would like to hear your thoughts. See questions.\n\nMinor things:\nOne detail that confused me a bit where the different start and end points of the x-axes in the plots. Maybe highlight that in the text or visually? It makes sense to separate it as you did.\n\nQUESTIONS:\nUnfortunately, to give a higher score I feel that the paper should go in a slightly different direction. \nQuestions for a higher score:\n\nWhich architecture modifications lead to faster/slower generalization? #Heads, Layers? Do these modifications scale linearly with the number of teachers or is another scaling observed? It seems that dMMSE requires non-linearity possibly provided by the softmax in the attention layer. What happens if you train linear self-attention Transformers on your settings. Do you generalize faster?\n\n\nOpen ended:\n1) Did you check other task diversities? It could be interesting to investigate e.g. different \"modalities\"  i.e. where the teachers are sampled from vastly K different curvatures? Maybe there are other distribution changes that make sense. \nThe goal would be to understand if one can push towards the generalizing solution with the same amount of data. \n\n2) Follow up on 1): Regarding your hypotheses of implicit regularization. It would be very interesting imo to check different initializations that usually are key to change between lazy and feature learning. It would be great to see if that intuition checks out or if there are other forces\nat play here. \n\n3) Can you plot the first two attention heads of the layer 1 and 2 (maybe more) for me for a model that is before in the middle and after the phase transition. My claim (maybe obvious) is that there will be visible differences here and that solutions that are close to ridge regression show copying behavior in the first layer to group together to do GD (see e.g. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html, https://arxiv.org/pdf/2212.07677.pdf) which might not be necessary to be implemented for dMMSE.\n\nLIMITATIONS:\nLimitations are obvious and mentioned in the discussion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper analyzes in-context learning capability of auto-regressive (continuous) Transformer models where the training data consists of linear regression problems. Paper finds that the ability to ICL with optimal solutions for the given prior (e.g. ridge regression for noisy linear regression) emerges only after the model is exposed to a certain number of different regression tasks (or simply a certain number of $w$s) from the prior distribution. After this threshold has been passed, however many instances they provided, the model choses Ridge regression instead of local Bayes optimal solution for the particular training dataset (called dMMSE).\n\nSTRENGTHS:\n1. Finds a threshold for task diversity that makes the Transformer model switch from local bayesian solution (dMMSE) to optimal bayesian solution (Ridge regression).   \n\n2. Find that weight norm improves this diversity threshold in the particular setup.  \n\n3. A clear and concise writing.\n\nWEAKNESSES:\n1. The results show this threshold exists even if they increase the number of instances per task, and number of training iterations. We do not know if the same transitions hold if the authors increase the number of instances per task and number of training iterations at the same time? Can you try an experiment where you maxed out everything: model size, training iterations and number of instances per task, and verify the results?   \n\n2. If 1 holds, what is the justification or intuition for this transition? As far as I saw, there was no discussion of possible theories to explain this behaviour.\n\n3. The experiments are only in linear regression with Gaussian priors, and previous work (e.g. Garg et al. 2022) had setups with sparse regression and non-linear problems.   \n\n4. Would be more complete if analysis of hyper parameters can be extended to the number of layers too.\n\nQUESTIONS:\n1. Are these transitions hold for different distributions than Guassian (Laplace, etc)?\n\n2. L75-79: Could you simplify the message here by giving the scaling laws (linear, logarithmic, or by its shape) of threshold given the task dimension. You can get this by running with more dimensions, for example, 2 to 128 dimensions and plotting the point where $\\Delta$dMMSE crosses $\\Delta$PT.  We can do same for weight decay and embedding experiments too. In my opinion, this type of plot would be much more explanatory than what Figure 5 and 6 currently has. And I do not think the different batch sizes are needed in these plots, why do not pick the highest one?\n\n3. Fig 6: Do we need 4 plots for different values of weight decay/embedding? Could you plot them at the same figure? Same for Fig5 left three panels. Also see my suggestions above.\n\n**Paper Summary**\nI think this paper finds something surprising that the models transition from dataset specific optimal solution to global prior optimal solution. But the results needs some verification and explanation as I queried in above.\n\nLIMITATIONS:\nThe authors could discuss the possibility of this transition might not be hapenning in the case of enough iteration and samples per task.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies whether a Transformer can in-context learn linear functions. While previous work studies a Gaussian prior over the weights, this work studies a discrete prior over the weights, resembling a less diverse pre-training task distribution. The Bayes-optimal estimator for this distribution is not Ridge Regression. This allows one to measure generalization based on whether the Transformer behaves like the Bayes-optimal solution to the pre-training distribution or the more general Ridge Regression estimator. It is seen that under a certain task diversity threshold, the model learns the Bayes-optimal solution to the pre-training distribution. Above this threshold, the model seems to learn the generalizing solution of Ridge Regression, indicating that pre-training task diversity may play an interesting role in ICL generalization.\n\nSTRENGTHS:\n1) The paper discusses an experimental setup that allows for the investigation of how task diversity can affect ICL performance, as well as the optimal solution.\n2) The paper varies batch size and step count, and both share a similar task diversity threshold where more training gets closer to ridge regression estimator. \n3) The paper assesses an optimally-tuned Smoothed dMMSE to eliminate a hypothesis for what the model implements\n\nWEAKNESSES:\nThough the paper identifies the existence of a threshold, there is not a strong study into possible reasons for why. There is a strong possibility that the threshold is a consequence of the evaluation setup rather than a feature inherent to the problem, and it is important for this paper to address these concerns since it claims that the threshold exists. I believe explaining the features of the threshold or why it exists would constitute a much larger contribution. Concrete suggestions below.\n1) One natural hypothesis, as noted by the authors on lines 235-236, is that the model simply doesn't have enough capacity to memorize M tasks. However, one could train a larger (or smaller model) while controlling for all other settings and seeing whether it has a different task diversity threshold. If it does, then this paper could offer insight into the features of such a threshold rather than just showing that it exists.\n2) It is possible that the threshold is a consequence of plotting the loss curves without giving more training steps to higher task count problems (the solution to 64 tasks is likely harder and takes more steps than the solution to 32 tasks). Is it possible to show the distance to Ridge over training time to rule out this possibility? Moreover, are there reasonable normalizations (alluded to in Question 2)?\n\nQUESTIONS:\n1) Is it the case that the model's distance to Ridge monotonically decreases over training or does the distance to Ridge increase and decrease? This can have important implications for the conclusion and determining whether the models for higher task count simply need to be trained longer, since it is possible the exact step count presented is still in the process of getting closer to Ridge and would later go further away.\n2) For the models that converged to sMMSE, how many steps does it take? Is it the same number of steps across task counts?\n3) (Minor) In Figure 2, is the upper right figure the same as the lower right figure?\n\nLIMITATIONS:\nThe authors do not have a separate limitations section but discuss alternative explanations for experiments throughout the paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}