Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression
Authors: Allan Raventos
Mansheej Paul
Feng Chen
Surya Ganguli
Venue: NeurIPS 2023 poster (2023)
Keywords: in-context learning
Bayesian inference
transformers
task diversity
emergence
TL;DR: We empirically demonstrate a task diversity threshold for the emergence of in-context learning in pretrained transformers beyond which the model can learn fundamentally new tasks in-context.

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
Allan Raventós"
Mansheej Paul" Feng Chen
Surya Ganguli
Stanford University
faravento, mansheej, fenge, sganguli)
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 updat- ing any weights. This raises foundational question: can ICL solve fundamentally new tasks that are very different from those seen during pretraining? To probe this question, examine ICL's performance 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 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 de- viating from the Bayes optimal estimator with the pretraining distribution the prior. This study also explores the effect of regularization, model capacity and task structure and underscores,
concrete example, the critical role of task diversity, alongside data and model scale, in the emergence of ICL.
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 []. This ability. called in-context learning (ICL), has unlocked the widspread use of language models by making it efficient to adapt general purpose models 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 architecture or training process; instead it emerges from pretraining on large-scale 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 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 similarity to tasks implicitly learned during pretraining. However, the distribution of tasks in our pretraining data, Teremnnn is usually
"Equal Contribution. Code released at http://ggtuubbcmmm/sassseeee///a task diversity
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

limited and unrepresentative subsample of the ideal distribution of tasks, Thre. that we want our model be capable of learning in-context. For instance, True could be the set of all instructions we want an A.I. assistant to follow. But, large-scale language modeling datasets 4 S
used to pretrain these models contain very few examples correctly formatted for ICI.. Instruction finetuning (IFT) datasets [6 designed to ameliorate this are expensive to collect and thus contain tasks from just few domains, Under the Bayesian framework, this distribution mismatch would cause the Bayesian estimator with a prior over the limited pretraining tasks, Trretrin 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 Trretrin and the vague specification of Thue make this question challenging to analyze. So. following recent work [4) we study ICL for linear regression. Here, a task is linear regression problem with 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 [3] 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, TPretminn 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, 7iree. 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 Trrctrain 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 as well as the optimal estimator for all tasks the Bayesian estimator with prior Trree- This allows us 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 Treetnin and perform suboptimally on tasks from Trae. or does it align with the optimal estimator for Time which can solve new, unseen tasks?
Contributions. Our contributions are as follows:
We find that transformer pretrained data with low task diversity behaves like the Bayesian estimator with prior :rremmin: 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 Trre-
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 sharp algorithmic phase transition in the limit of many examples per task, aligning with the optimal estimators on Trretrin before the threshold and on Tire 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 at 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 Terctrii grows progressively less similar to the optimal estimator for Tire-
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 ICI.
Overall these contributions suggest that the emergence of in-context learning in pretrained transform- cannot be fully explained by theory of Bayesian inference on the pretraining distribution.
2 Problem setup
ICL of linear regression (schematic in Fig. D.. Each ICL lask corresponds to a latent D-dimensional regression vector, w E RD. At inference time, the transformer takes as input a sequence of K data- target pairs, (X1, .i.. XK: yk), which are the in-context examples corresponding to this single task
2

---

w
(o
fo
fo(Sx)
Xki Vk
Transformer
fo
o
X2, V/2 9
1
1
1
X1
11
X2
1/2
Xx
91
$2
Figure 1: Schematic for ICL of linear regression. (Left) A task corresponds to latent regression vector, w (purple line). (X1, (11.31.... 0i..., XK 8/K (black circles) is sequence of in-context examples for this task. (Right) The PT, fo, takes this as input and generates K outputs. The kth output, (o( is the prediction for the target of xg and depends only on the context Ss = (X1: y1,.., Xk-1, 11 Xx)-
w. For & E {1,..., K}, the data are drawn i.i.d. from a D-dimensional standard normal distribution, X2 ~ (0, Ip), and the targets are scalars given by yı w X₁: T drawn i.i.d. from a normal distribution with mean and variance 0² Ex NO, o Let fo be the PT with parameters 0. Since use decoder-only transformer with a causal attention mask, for each & E (I,..., K}, the transformer sees the context S = (x1, 31, - 1k-1, Xx) and based on this context, it makes prediction, fo (Sk). for the target of X 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.
&s. The Ex: 's are noise scalars
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, Trretrain This distribution has limited diversity as the uniform distribution over a finite of M tasks, Trerrain = U( (((w(1) ..... w (MM). Each task in Trretrii drawn i.i.d from D-dimensional standard normal distribution, (()) (O, ID), i E 1,..., M. By increasing the number of tasks, M, in Trrevain. 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 Themis is just the MSE for each prediction, averaged over the predictions in the sequence:
L Peemin (0)
1 K
(So(Se) -ya)²
W Theepen (0,T)
M (00
Evaluation. We evaluate the PT's performance on tasks seen during pretraining by computing L Theeis 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 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, Tirees over all latent regression vectors; in our case Tiue = A0, Ip). We evaluate the PT's performance on new tasks by computing T.Time, which follows Eq. but where the tasks are sampled from the ideal task distribution:
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" in Eq. ( for both task distributions, Treuran and Tiree- The optimal estimator for the kth prediction, oZ that minimizes the kth in the sum in LT is the Bayesian estimator with T as the prior. This is given by the posterior mean of yu conditioned on the context: o% IT. [b% Sk], where the expectation is over the task distribution, T, and the noise, Ek (Appendix A.I)
For task distribution Trretrin (1)(1)) ...,, ((M) }. the discrete minimum mean squared error (dMMSE) estimator is optimal. It is given by JOMMSE (WMMMS)) X where wMMMS

---

0.4
10
10-2
Trretrain
MSE/D
0.3
APT. aMASE
APT, Ridge
0.2
10-3
10
0.1
10-5
10
1.5
Ridge dMMSE
10-1
Batch Size
10"2
256
MSEID
PT
APT. awASS
APT, Ridge
512
Trree
1.0
10
Batch Size
1024
256
10-3
512
0.5
1024
909099000000
10-4
2
2
212
2
220
20
2
2
26
2
220
Pretraining Tasks
Pretraining Tasks
# Pretraining Tasks
Figure 2: ICL emerges in PTs beyond threshold pretraining task diversity. We show all results on both tasks seen during pretraining (lop row) and on new tasks (boltom row). The left column compares the normalized loss of transformers pretrained with increasing task diversity that of dMMSE and Ridge. When the pretraining task diversity is small, the PT's performance matches that of dMMSE; performs very well on tasks seen during pretraining but poorly on new tasks. As the pretraining task diversity increases, both dMMSE and PT approach Ridge. However, the PT approaches Ridge much faster, significantly outperforming dMMSE on new tasks (bottom left). In the middle and right columns, we compare the PT's predictions to those of dMMSE and Ridge respectively (Eq. 44) We also increase the number of sequences per task each level of task diversity by increasing the batch size while keeping total training steps fixed. This reveals task diversity threshold between 214 and 215 pretraining tasks which there a phase transition behavior of the model. Below the threshold, increasing the dataset size leads to PTs with predictions more aligned with dMMSE on Tpretrain (top middle). However, beyond this threshold (indicated by gray shading), increasing the dataset size leads to PTs more aligned with Ridge on all tasks (right).
1 EM w and for k (2,... KJ. (Appendix A.2 M
MMME
M
exp
20
j=1 (y) w (xXxx))
w (())
(2)
E)=1 w()Txx)2
1=1 exp
20
Intuitively, W.MMSF is just a weighted sum the pretraining w(0s with weight governed by the likelihood of observing targets (y1,-.., M/-11) conditioned inputs (x1,-.., Xx-1) } and the task being w(), A PT that minimizes the pretraining loss I, Thamii will behave like this estimator.
For task distribution True (O ID), the Ridge regression estimator with the ridge parameter set Ridge) the noise scale r² optimal: y (w
Ridge
Xt. where w
and for k {2, ....... K}.
Ridge
(X + I'I X
where X X= =x]x, [T...,,[__1)) (x,,x,,,,xX__) R(K-1)xD and y (y/1-..,, .k....--i1 (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 given task distribution T. We write this as
T AM.R
1 KD
(o((k))
RidgooUMMEE
)
6 (00
3 Experiments and results
Unless specified otherwise, we study linear regression D) dimensions with up to K 16 in-context examples and observation noise variance o² 0.25. We use either base transformer

---

Tpretrain
Trrue
Teretrain
Truee
10
10-1
10-2
10-2
APT, aMMSE
APT. awoesE
ApT. Ridge
BPT. midge
10-3
10-3
Steps
10-3
10-2
0.5M
10-5
10-5
1M
10
10
2
2 2
2
2
2
2
2
2
2
2
2
2
2 2
2 2
2 2 2
Pretraining Tasks
# Pretraining Tasks
Pretraining Tasks
# Pretraining Tasks
Figure 3: Incrraaeed pretraining steps reveals the same task diversity threshold for the emergence of ICL. Columns I and 2 in this figure are similar to the middle column in Fig. Big. columns 3 and 4 correspond the right column in Fig. 2 except here we increase the number of sequences per task by increasing the number of training steps while keeping batch size 256. Both methods of increasing dataset size- increasing batch size Fig. 2
and increasing training steps in this figure-reveal a transition in the behavior of the PT: beyond the task diversity threshold, ICL on new tasks emerges.
model with the GPT2 architecture 15 with 8 layers, 128-dimensional embeddings, and 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 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, Trrcuainn described in Section 2 We then minimize the objective L Trecin in Eq. ① using minibatch stochastic gradient descent. For each sequence in a minibatch, we sample single task w from Trretzinn well as new samples of data, {x;
and noise, {s,) from their respective continuous distributions, to form sequence (x1, [x1 + El. XK, w + EK). If we train for N steps at batch size B, the transformer will see total of NB unique sequences and roughly NE unique sequences for each latent task in Trremin- By increasing either or N at fixed "M,
M, can increase the total size of the pretraining dataset (or number of sequences per task) while keeping the dataset diversily-the number of unique ws in Pretrain fixed.
3.1 Task diversity threshold for the emergence of in-context learning
For F 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, = 500K). We evaluate the PTs and both optimal estimators on tasks seen during pretraining drawn from Trremmin (Fig. top left) and on new tasks drawn from True (Fig. 2 bottoo left) and plot MSE normalized by task dimension-22 -T /DD from Eq. (D). Since dMMSE optimal tasks from Trretrain (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 L'Taani explicitly encourages the PT to match dMMSE performance. On the other hand, Ridge is optimal on tasks sampled from Tiree (Fig. 2 bottom left); the blue markers denote the lowest possible MSE the PT could attain 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 diversity-M M up to about 2 2te PT's MSE closely tracks that of dMMSE on tasks sampled from Trettrin (Fig. top left); the PT performs optimally on tasks seen during pretraining. But it significantly underperformy
new tasks sampled from Tirec- indicated by the gap in MSE between the PT and Ridge (Fig. 2 bottom left). In this regime,
behaves like the Bayesian estimator with prior Teretriin
High task diversity phase: the PT is non- non-Baaesian with respect to the pretraining task distribu- tion and can solve new tasks. At higher task diversities- aboveaabvveeee 2 pretraining tasks-the PT's MSE deviates from dMMSE and approaches Ridge under both Trretrin and Tirue- Crucially, the PT starts to significantly outperform dMMSE on unseen tasks sampled from Tizee
bottom left) at the expense of not fully minimizing its training objective, Themis (gap between PT and dMMSE under Threrriin Fig. top left). This suggests that, PT trained on 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.
5

---

Small PT Learning Curves: BS 512, 500K and 2M Steps
Scaling of Early Stopping Time
DpT. Roge on Frrue
0.08
1e6
0.06
2 1e5
0.04
t
0.02
to 0.47)
1e4
2e6
250K
500K
1M
2M
2
Step
Step
Figure 4: Learning dynamics of small PTs shows a transition at the task diversity threshold. We plot A Thamis vs training steps for small PTs. For the same M, learning curves for short (500K steps, PT.Ridge
left) or long (2M steps, center) training durations are similar, and for M M = 211.5 learning curves are similar to that of a model trained with infinite task diversity. Right: For M < 200, t* (the
training step at which A Thamin is minimized) is well modeled by scaling law t" x M A linear fit Pr.Rdde
of log 1" vs log M (dashed red line) gives a R 0.47. But for M > 210, A Taman decreases through training; = 2M, is larger than predicted by the scaling law. This sudden break in the scaling law suggests fundamental difference in the learning dynamics of models on either side of the threshold.
Finite size scaling of training data suggests an algorithmic phase transition as task-diversio increases. The experiments in Fig. (left column) suggest that, when tested on both task distributions Trretrin and Trre.. the ICL algorithm implemented by a PT exhibits a smooth crossover performance from dMMSE to Ridge. We next examine how this transition changes as we increase the number of sequences per task seen over pretraining, the transformer sees more sequences per latent task in Trremrinn 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 manner that facilitates ICL on Time-
fixed task diversity. One might reasonably expect that,
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 PT. which quantifies how different the PT and dMMSE estimator's 's predictions are when testing on tasks drawn from Trrttmin does in fact decrease for M < 210 (Fig. 2 top center) as we train on more sequences per task. Moreover, for each M E {210 ..... 2113 the PT's predictions also become less similar to those of Ridge, both on tasks from Trectran (Fig. 2 top right) and Time (Fig. 2 bottom right). Crucially, this movement in behavior of the PT towards dMMSE and away from Ridge, at least on tasks drawn from Trrurain holds only up to a threshold number of tasks between 214 and 2 Beyond this threshold, pretraining on more sequences per task at a fixed task diversity actually makes the PT more like Ridge, in that both A Thamie PT.Rigg and PT.idge 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 Time 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 Teretain at task diversities beyond the threshold.
A Thamis
True
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 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 moderate task diversity threshold below 2 pretraining tasks; even though dMMSE significantly underperforms relative to Ridge on Tieee 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 I 1M while keeping batch size fixed at 256. We observe that doubling N (change from pale blue to red in Fig, and doubling B (change from pale

---

2 Pretraining Tasks
2 Pretraining Tasks
2 Pretraining Tasks
0.6
0.5
MSEID
Ridge
dMMSE
0.3
PT
0.2
0.0
0.2
0.4
0.6
0.3
1.0
0.0
0.2 0.4 0.6 0.8 1.0
0.0
0.2 0.4 0.6
0.8
1.0
Interpolation
Interpolation
Interpolation
Figure 5: Transformers pretrained with high, but not low, task diversity can learn new tasks in-context. We compare the normalized loss of the PT to that of dMMSE and Ridge as we interpolate between tasks in the pretraining dataset. Left: At 2 tasks, well below the task diversity threshold, the PT performance matches that of the dMMSE estimator along interpolating paths, but under-performe Ridge on new tasks at the center. Middle: At 20 tasks, the PT outperforms dMMSE on new tasks at the center of the interpolation path, but is not yet good as Ridge on new tasks. Right: At M = 2 tasks, just above the task diversity threshold, the PT performs as well as Ridge even on new tasks at the center. This demonstrates that, when pretrained on data with finite but large number of unique tasks, the PT, unlike the Bayes optimal estimator for Trreuamn can learn new tasks in-context.
blue to red in Fig. 2) have very similar effects on ATT.MMSE and FI for both T = Trre and
T == Trretain- More importantly, the task diversity threshold, which we determined as the cross-over point in A Time
Ri.Ride 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 IM steps at batch size 256 (Fig. [] right). Given that 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 consider the task diversity threshold to be cross point in A Thee RidRi between batch sizes 256 and 512 when training for 500K steps. See Appendix 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 each number of tasks.
Learning dynamics and a break in the sealing 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
algorithmic phase transition but at lower task diversity threshold between 2 and 2 pretraining tasks. In Fig. [eet, we visualize the learning curves (Aprrsdge 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*, PT.Ridge decreases early in training until reaches a minimum at time step t", and then increases as the PT approaches dMMSE. We define 1* as the early stopping time for Ridge. For M M*, PTRidde decreases throughout training. To evaluate if, the latter case, models are undertrained and t" larger than the total training time, we extend training to 2M steps batch size 512 (4x the training time, see Appendix B). Fig. Acenter, 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. TO), suggesting that these models are approaching the Ridge solution.
A Time
A Time
In Fig. 4 right, we study how t", scales with M. For most M M t" obeys simple scaling behavior t" x _M,, a == 0.47. However, for M > 20, the distance to Ridge decreases monotonically through training and 1* = 2M steps. Despite the caveat that our experiments are necessarily in the large but finite training step regime with decayed learning rate schedule, this stark break in the scaling behavior of 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.
7

---

8 Dimensions
16 Dimensions
32 Dimensions
M = 2
10-1
Batch Size
10-1
10-1
0.5
256
DMedel Ridge
0.4
Trree
bpp, halge
10-2-
512
10-2
10-2
0.3
Model
10-3
10
10-3
0.2
dMMSE
0.1
10-4
10-4
10-6
0.0
2
2
2
2
220
2/2
2
2
2 2
2
2
2
2 2
16
24
32
# Pretraining tasks
# Pretraining tasks
# Pretraining tasks
Task Dimension
Figure The task diversity threshold increases with task dimension, and the PT's ability to solve new tasks scales significantly better than dMMSE's. We vary the dimension of the regression problem (first three panels) while leaving the signal-to-noise ratio fixed. The task diversity threshold consistently increases with task dimension (gray shading denotes post threshold). At 22 tasks (right), the PT's predictions are similar to those of Ridge at all D (orange A Time PT.Ridge)
whereas dMMSE grows progressively less similar to Ridge (blue A TTme Radge.aMMSE.
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, pairs of seen tasks w, and w, in the support of TPretaiin The interpolation path is given by
one parameter family of new tasks w. that interpolate between
1 (11 (1wall2+ + l|w)ll)) llaw; (1 (1-a) Wa 2
ow; + (1 (1-aa) a)),
for E [0,
(5)
Here we fix the norm of the interpolated vector Wa to the average of the two endpoint norms avoid ||wall |I taking on very small values for a ~ 2. Fig. 5 shows the results of this analysis for 2 (left, low task diversity regime), 20 (center, below task diversity threshold), and 215 (right, just above the task diversity threshold) tasks. At each value of a, MSE is averaged over large number of task pairs (W;, wj). 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 smoothed dMMSE model. We have seen that at an intermediate task diversity the PT significantly outperforms dMMSE on new tasks in Trme- It is clear why dMMSE performs poorly on new tasks in Trree at low task diversity: its prior over tasks concentrates on M unique tasks in Treetrinn while the prior over tasks in Tiue is Gaussian. A natural conjecture is that the PT cannot memorize all M tasks in Trretrain 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 Time (see Appendix G for details). This smoothed dMMSE outperforms dMMSE as it has a prior over tasks closer to the Gaussian Trre 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 2D, and increasing observation noise x² to match the SNR to that of D = 8 and x² = 0.25. We also train a larger model with 12 layers, 256-dimensional embeddings, and 4 attention heads that 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 24,, 1 and 2 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 Time to sphere for large D. Thus we expect dMMSE performance to scale much more poorly with dimension D since the finite number of tasks Terettain would need to cover substantial portion of the sphere for dMMSE to approach Ridge. To test this hypothesis, for M 220 which is
8

---

Ablate Weight Decay
Ablate Embedd ing Dimension
Ablate # Layers 214
Tosl Diversity Threshald
2
Task Task Diversity Threshold
2
Task Diversity Threshald
2
21
243
213
222
222
Base Small
212
2
222
210
211
211
0.01
64
128
256
512
1024
6
10
Weight Decay
Embedding Dimension
Layers
Figure 7: Explicit regularization and model capacity affect the task diversity threshold. Increas- ing explicit regularization, in the form of weight decay, lowers the task diversity threshold in base PTs (left). Increasing embedding dimension, has no effect on the threshold for base PTs, but does increase the threshold for a small PT (middle). Increasing depth, while holding other hyperparameters fixed, increases the threshold for a small PT (right). See Figure 13 for plots of A Tine
Rids vs M.
the largest task diversity consider, we explore how the similarity of PT and dMMSE predictions to Ridge new tasks scales with D (Fig. 6 right panel). We see that A Time Ridg grows significantly
as we increase D, while remarkably A Fiue RidR 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 all D for M 2 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. 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. (rggter) or depth (Fig. rightt increases the task diversity threshold. Base PTT has much larger capacity then small PT and also has larger threshold; we hypothesize that small PT still in regime where the threshold is sensitive to capacity while buse PT not. Together, these results suggest that model capacity plays an important role in the emergence of in-context learning: increasing capacity (up to point) leads to an increase in the task-diversity threshold.
4 Related work
The Bayesian framework for ICL introduced by Xie 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 1& 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 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 fully explained by Bayesian inference.
Prior work 12 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 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 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 Tremman and True. This allows us to establish that ICI, 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. [2S] study how ICL in large language models is affected by semantic priors and input-label mappings, focusing differences across model scale. Olsson et al. [29] study inductions heads- circuits responsible for completing patterns by copying
tokens
mechanism for implementing ICL.
5 Discussion
Overall, we have extensively explored the impact of pretraining task diversity the emergence of in- context 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 Trrmmin (i.e. dMMSE), or for the diverse ideal task distribution TTree (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 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 break free of the pretraining task distribution to solve new tasks, given moderate pretraining task diversity.
Remarkably, beyond the task diversity threshold, PT's 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, PT's 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) [1]. In Fig. [[]e we find that PTss trained with large M inhabit the same loss basin as PTs trained with M = oc: the training loss barrier between PT's trained with M 2 213 and PT's with M = oc is similar to two PT's trained with M = co. In contrast, there are large loss barriers between PT's trained with M < 213 and M = cc. Additionally, PTss trained with M = oo are closer in weight space to PT's 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 Trmei 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 categorize the contents of such corpora according to tasks and measure their resulting task diversity. Regardless, our observation in linear regression that moderate threshold in pretraining task diversity can enable PT's 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.
10

---

Acknowledgements
The authors would like to thank the Google TPU Research Cloud (TRC) whose generous support made this project possible. SG thanks NSF CAREER award and NTT Research for support.
References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child. Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz 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 http:pps://pocceeeeddteee neuris cc/ wacc/pppr__a f Files/ wwwlss/pappeeee peless/apperr222///iiee.. Paper. pdf
[2] Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant, Risks from learned optimization in advanced machine learning systems, 2021.
[3] 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.
[4] 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 unified text-to-text transformer. J. Mach. Learn. Res., 21(1), jan 2020. ISSN 1532-4435.
[5] 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.0000 2020.
[6] 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:///pennrevier htt/fforummid-gg22GGoom
[7] 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 Szezechla, 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. Multi- task prompted training enables zero-shot task generalization. In International Conference on Learning Representations, 2022. URL httpss//opoeeeeiieess net
[8]
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasckaran, 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- Naarallttttlnnssttttttonnns: 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- wwa-maann... 340
11

---

[9] Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL: 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 Linguisties. doi: 10. 18653/11/ 2022.naacl- 1022.naallmmin.201..
org/ 2022. naacl-maii 201
[10] 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: //pooe/eee///oooeeenngg.gg/ wwaeinngg..neuiiiee www. wwtppppppppppcpccccccc http://aarrrrreeeeeeeeeeeeeee files/papeee Conferene pdf
[11] 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 13388, 2023.
[12] 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. URLhttps:: //wpneveeeeeeeeeee////ee/ee wttppperww...eeeeeeeeoaraa//d//a ret.:forrummiit-f122222eee
[13] 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 Inter- national Conference Learning Representations, 2023. URL http:: httLhtts://ppprrrreeie http::::///pereiee... net/
forum? Dorum?id=OgOx4HByYl
[14] Johannes von Oswald, Fyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2022.
[15] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Iyya Sutskever. Language models are unsupervised multitask learners. 2019.
[16] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CORR, abs/1412.6980, 2014.
[17] Leslie N. Smith and Nicholay Topin. Super-convergence: very fast training of neural networks using large learning rates. In Defense Commercial Sensing, 2017.
[18) Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in- context 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.naacI main. meee2nnacl-mmin1111 URL httpp:://aaaattaaaaa
202g/22222.gg www.l-maiiia 191
[19] 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 In- side Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pages 100-114, Dublin, Ireland and Online, May 2022, As- sociation for Computational Linguistics. doi: 10.18653/1120022.deeli 1.10. URL https //aclanthology.co org 2022. deelio-1.
[20] 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.
[21] 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 1 11048-11064 Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. htRpst http:/:::///aaaahoolog/// org/ www.o02..........oo. 2022. emnlp emnlp-main 759.
12

---

[22] 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 meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023. URL /nttps://openreiieg httt::///ppprrrreeeeeeee/nnnadan rzo
[23] Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Meta- learning without memorization. 2020.
[24] Ramnath Kumar, Tristan Deleu, and Yoshua Bengio. The effect of diversity in meta-learning. 2022.
[25] Nilesh Tripuraneni, Chi Jin, and Michael I. Jordan. Provable meta-learning of linear representa- tions, 2022.
[26] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in- context learning by meta-learning transformers, 2022.
[27] Stephanic 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.
[28] 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.
[29]
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Danny Hernandez, Scott Johnston, Andy Jones, Jackson
Deep Ganguli, Zac Hatfield- 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. heads/index.hrr html.
[30] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis, 2020.
[31] Mansheej Paul, Feng Chen, Brett W. Larsen, Jonathan Frankle, Surya Ganguli, and Gintare Karolina Dziugaite. Unmasking the lottery ticket hypothesis: What's encoded in winning ticket's mask? In The Eleventh International Conference on Learning Representa- tions, 2023, URL https: https://openreett //pp:/eveee/eeeeee///me www../oorrmmdd--XSWWaa.... ukZ
13

---

A Bayesian estimators
A.1 Bayesian MMSE estimator
The optimal estimator for the kth prediction, 9% that minimizes the kth term in the sum in L (in Eq. (T)) is the Bayesian estimator with T the prior. This given by the posterior mean of ye conditioned on the context: 9% = ET 3ik Sk]. where the expectation is over the task distribution, the [fo(S))) yua)² |sa)
T, and the noise, Ek. To derive this, notice that, using the law of total expectation, the kth loss L is:
E
(fo(Se)
yw)
E WhT =x~N(0.0²)
Se
which implics that the optimal estimator is o7 Ex [y Sy]- Adopting the notation X (xT...., xx_1) R(k-1jxD and y (y1,..., Uk-1), we express the expectation explicitly:
f dwdyk ykp(w, yh X, y, x;x)
E Luk Sk]
/ dwdyk yaPlyk xxs w)p(w x,
/ dw T p(w x)
= wIxE
where
S (dwp(w) (oyp(w)ww p(ya (Xi, w) dwpp p( X;, w)
W E w Sk]
(6)
A.2
dMMSE estimator
For task distribution Trretrain U(w(() (ad)), we can directly get the discrete minimum mean squared error (dMMSE) estimator by plugging the uniform discrete distribution into Eq. 6
dMMSE
M
exp
2
1 Eイ wx))xx
w.())
M Ei=1
(-22- E w()Tx,)
exp
A.3
Ridge estimator
For task distribution Time = N(0, Ip), we can directly get the Ridge regression estimator from
Eq. 6
22 Xww - (Xw y) w w w]dw (Xw 7) (Xw y) w w dw
Ridge
w exp
f Texp
wexp
- (w (x x 111)) xx L
-xx) dw
ʃ exp
7 (w (XTX IlIn) -Xy) (x
(XTX - dw o o2I
[XTX +o I²ID) -xx,y
(8)
where X (x], x 1(-D and y (y/..... 1k-
B Experimental details
For most experiments, we study linear regression in D = 8 dimensions with up K = 16 in-context examples and observation noise variance r² 0.25. Our base model is transformer with the GPT2 architecture [15] with 8 layers, 128-dimensional embeddings, and attention heads. We train with
14

---

the Adam optimizer [16] and one-cycle triangle learning rate schedule 17 with 50% warmup. The base model trained with batch size 256 for 500K training steps. though these hyperparameters are varied in our experiments. Specifically, for the experiments in Fig. 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 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 32 setting. We also train the models on D = 8, 16, 24, 32 with up to K = 2D in-context examples in each case. The observation noise variance is scaled at each dimension keep the signal to noise ratio constant. So x² = 0.5, 0. 00707, 0.866, 1 for D = 8, 16, 24, 32 respectively.
We always sweep at which training is stable. For almost all experiments, the learning rate 0.001, the main exception being the experiments in Fig. 6 where the learning rate is 0.0001.
learning rates in (0.0001, 0.0003, 0.001) and choose the largest learning rate
For the 2M training step experiments Fig. 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 at 2M steps, consistent with our choice of decayed learning rate schedules. The experiments in Appendix FFaee 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 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. S provides an additional visualization for Fig. 2 uupper middle but on a linear scale on the y-axis. In this figure it is clear that for all task diversities below the threshold- that increeainn the batch size aligns the PT with dMMSE, the optimal estimator for Treeurin Thus below the task diversity threshold, the optimal estimator does indeed approach the Bayesian estimator with prior Trerrain when trained on more data and is thus unable to learn new tasks as discussed in Section3.11
Batch Size
0.08
256
0.06
1024
0.04
0.02
0.00
2
20 2
Pretraining Tasks
Figure 8: Another visualization for Fig. 2 upper middle but with the x-axis extending all the way to the task diversity threshold and a linear scale on the y-axis to make it easy to differentiate the different batch sizes at large number pretraining tasks. See Appendix for details.
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. middle and right column, we probe this by increasing the batch size to increase the number 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.
15

---

In Fig. [9w 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 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 kept constant, atleast within reasonable range 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 5 512 and 1M training steps for batch size 128. In all three cases, the behavior of the PT identical as measured by A ATT.MMSE and A PT.Rigee for both Trretrin and True across all numbers of pretraining tasks.
Tpretrain
Latent
Teretrain
Latent
10-1_
10-1
DrPo. dMMSE
ДPTM. dMANSE
Arna, Ridge
10
Drna midge
10
Batch Size, Steps
10
10
--- 512, 250K
10-3
10-3
256, 500K
10-5
10-5
128, 1M
10 4
10 4
2
2
2
2
2 2 2a 2
2 2
20
2
2 2 2 2
Pretraining Tasks
Pretraining Tasks
Pretraining Tasks
Pretraining Tasks
Figure Number of sequences per task and not batch size or number of training steps affects final model behavior: This figure provides additional experiments to support the findings in Section 3.1 and Figs. 2 and Each panel corresponds same panel Fig. See Appendix [ for details.
E Small PT
We run our task diversity threshold finding experiment a small PT with 64 dimensional embeddings, 4 layers, and 2 attention heads. The existence of 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 lower capacity has lower threshold (~ 211.5 compared to ~ 214.5 for the Base PT Fig. 2 but also has lower overall performance
(compare y-axis range to Fig. 2 lower right).
Small
Batch Size, Steps
256, 500K
APT, APT, Ridge on Tirue
512, 500K
10-2
512, 2M
M
10-3
20
22
21
26
28
(# Pretraining Tasks)
Figure 10: Existence of task diversity threshold reproduces in small PT. Starting from batch size 256 and 500K training steps, we train on 2x data (i.e. examples task) by doubling the batch size, and on 8x data by doubling batch size and quadrupling the number training steps. We show distance to Ridge trained with infinite task diversity, where each training example samples a new task from the Gaussian distribution, dashed lines (colors indicating batch size and training time). Points after the threshold (M 2 222) are shaded gray. Importantly, the threshold value doesn't change even if we go from to 8x data suggesting that it isn't 't artifact of underfitting.
new tasks (as in Fig. lower right). and also plot distance to Ridge for models
16

---

Tpretrain loss barrier
1.5
180
M
finite
L distance
M finite
1.0
160
140
0.5
120
0.0
100
10
15
10
15
logz (M)
log2 (M)
Figure 11: The training loss barrier between PT trained with M 2 21 and a PT trained with M = co is small. At each M, including M oc, we train four small PTs, each with different random seed for tasks. Each purple dot is an average over the finite-infinite task model pairs the given value of M, and each light blue dot is average over the 6 definite-infinite task model pairs. Error bars are standard deviations. See Appendix for details on how loss barriers and distances between models are computed.
F Examining the loss landscape along paths interpolating between finite and infinite task models
For each finite number tasks, M, as well as M do (where each training example samples a new task from the Gaussian distribution), we train four small PTs, cach with 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 = oc models, both in terms of the size of the training loss barrier along path interpolating between the models, well the distance in weight space between the models.
We compute the training loss barrier between finite M model and an infinite task model under the objective Tretrain defined by the specific set M tasks the fimite task model was trained on. We follow the approach outlined in (31] where the barrier between two models is the loss at model whose weights are the average of the two models, minus the average of the losses the two models. Specifically, given two weight configurations wh and wy and loss function L, the loss barrier is calculated as
+ W2 L(w1) L(wx)
loss barrier
The distance in weight space, in turn, is simply the L2 norm of the difference in model weights. For each M we average barriers and distances across finite-infinii task pairs of models. We treat the analogous quantities computed between pairs of infinite task models as baseline. See Fig. and the Discussion section in the main text.
and
G Smoothed dMMSE
Frre
Ridge
Smooth
MSE/D
0.346
0.345
212
24
2
218
220
Pretraining Tasks
Figure 12: The PT outperforms the Smoothed dMMSE estimator with optimal smoothing. See Appendix G
17

---

Here we consider the setting in which the discrete point prior over M tasks seen in pretraining is replaced with mixture of M isotropic Gaussians with variance e². We call the estimator that does posterior predictive inference under this prior the Smoothed dMMSE estimator (sMMSE). In the limit o O, the sMMSE estimator becomes the dMMSE estimator. On the other hand, if c - oc, the prior will look like a Gaussian distribution but will have a variance much larger than the variance of Time In both edge cases, it won't perform well on Time and so there is an optimal e that achieves the best performance on True. In practice, we perform search by simulation to determine the optimal We observe in Fig. 12 that for 2 tasks onward, the PT outperforms the optimally smoothed sMMSE on Tiree-
To derive the sMMSE estimator, we want to calculate Wa
w Sk]] and apply Eq. 6
p(S p(Sx(w)p(w)
p(w SAa
((5))
1-1 1 x
M
1
(w - m()))
p(y Xx,
W)
exp
2₂2
/=1
1=1
M
k-
(Txx - -v;)²
1
ww w(0)/2
exp
20 /=1
20
1=1
(1 1
02 (x2X + 2
1
W+ (1 will)
x2xly T
1
/w(lll1))
cxp
WT
2x2
We define,
1 XIX +
1
A
8²
J. 2
w(O
+
1
X
0²
M
(-1 1
2ea /1/()lpp) / dw exp ( (-2
1 2 WT SWAAW
== Σ exp
Aw + JJ.w)
1=1
(2T)
M
J.A A
-AA-JJJepp
Teexp (-1 1
2c /1m(i11))
A
t=1
so that,
M
( Ww Aw + J[w 2 ll (01p)p))
1
p(w
Sx) =
exp
C 1=1
Plugging back to Eq. 6
W t/
MMMS
M
exp
2e mm0012)
) dw /dw exp
WWW + +J[w)
C 1=1
1
(2x)) D
exp (F
1
A
J
C
A
2.2
M
3 BAA-'J;
where 3 = softmax(3) and B, = = -1J;
| e 12
2ca
H
Supporting figures
18

---

Weight Decay: 0.0
Welght Decay: 0.01
Weight Decay: 0.1
Welght Decay: 1.0
10-1
Batch Size
DPT, Ridge
10
Frrue
10-3
10
2
2
2 2 2 2
n 2 2 2 2 2
2 2 2 2 2 200
2 2 2 2 2 2
# Pretraining tasks
# Pretraining tasks
# Pretraining tasks
# Pretraining tasks
128 Embedding
256 Embedding
512 Embedding
1024 Embedding
10
Batch Size
Trrue APT, Ridge
10
2 214 2 2 220
2
2 2
2 2
2 2
2 2 2
2
2
2 2 2
# Pretraining tasks
Pretraining tasks # Pretraining tasks
Pretraining tasks
Figure 13: Effect of weight decay and embedding dimension on the task diversity threshold of base PT. Increasing explicit regularization, in the form of weight decay, consistently lowers the task diversity threshold, which the crossover point between Pi.Ridge at the two batch sizes (top. gray shading denotes post threshold). However, note that more weight decay does make the PT's predictions less similar those of Ridge, as quantified by A Tire Ridge numbers of tasks beyond the task diversity threshold (i.e. the flat region of both blue and red curves progressively move up with increasing weight decay from left to right in the). Increasing embedding dimension, while scaling the number of heads by the same factor, has discernible effect the threshold (bottom) for base PT.
Laplace prior
Batch Size
MSE/D on Frre
0.345
256
512
M
0.340
0.335
12
2
2 16
2
Pretraining tasks
Figure 14: Task diversity threshold for tasks sampled from a Laplace prior. We run our task diversity threshold finding experiment in setting where each element of each task regression vector sampled from a Laplace prior instead of a Gaussian prior. We use the base model from our paper. For the baseline, we approximate the Bayes-optimal model new tasks by training for 500K steps with batch size 512 on pretraining distribution with infinite task diversity drawn from the Laplace prior (black dashed line). The y-axis is the dimension normalized mean squared error on new tasks drawn from the Laplace prior. In this setting, we see task diversity threshold between 214 and 15 pretraining tasks; beyond the threshold (gray shading). training on more data leads to better performance on new, unseen tasks and ICL emerges. Thus, our findings for the task diversity threshold reproduce when tasks are drawn from different priors.
19


Reviews:
[{"review_id": "o5h5nWHVQF", "reviewer": "Reviewer_vcDA", "summary": "This 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.", "questions": "1) 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", "limitations": "The authors do not have a separate limitations section but discuss alternative explanations for experiments throughout the paper.", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1) 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", "weaknesses": "Though 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)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful questions which motivated new experiments which we hope are compelling:\n\n**Effect of model capacity on task diversity threshold (Weakness 1):** In the rebuttal, we further explore the effect of model capacity on the task diversity threshold. We use a Small PT with 64 dim embeddings, 4 layers and 2 attn heads. We hypothesize that this model is more capacity constrained than the base model in our paper (Base PT) and so ablating parameters affecting capacity will have a larger effect on the task diversity threshold. We observe that changing from Base PT to Small PT leads to a drop in the task diversity threshold by a factor of 8 from $2^{14.5}$ to $2^{11.5}$. Furthermore, if we start from Small PT and either increase the number of layers (Fig 1B), or embedding size (Fig 1C), the task diversity threshold goes up, demonstrating that, when starting with smaller models, model capacity can be a key feature of the task diversity transition. In our paper, we showed task dimension (Figure 5 main text) and weight decay (Figure 6 main text, top) could also control the task diversity threshold. We believe that training Base PT with weight decay decreases the effective capacity of the model, hence lowering the task diversity threshold. Together these results provide a more extensive characterization of how layers, embedding dimension, weight decay, and task dimension all affect the task diversity threshold, as requested by the reviewer. \n\n**Is the task diversity threshold an artifact of undertraining (Weakness 2, Question 1 + 2)?** We refer the reviewer to the discussion of undertraining in the global rebuttal. Here, we emphasize a few key points. Following weakness 2 + question 1, we plot distance to Ridge training curves for 500K and 2M steps (rebuttal Fig. 2A). At either training duration, training curves fall into 1 of 3 groups: (1) at very low task diversity (very small $M$, not shown in figure) models monotonically move further from Ridge, (2) at low task diversity (small $M$) models move closer to Ridge then move away, finding the dMMSE solution, (3) at high task diversity (large $M$) models monotonically move closer to Ridge. The transition from pattern (2) to pattern (3) happens near the task diversity threshold. Comparing Fig 2A left and right panels, we note that if one increases the training time even **by a factor of 4** from 500K to 2M steps, this dichotomy between the behaviors of training curves on either side of the task diversity threshold is preserved. Also, form Question 2 about normalization of training steps, it is indeed harder to learn dMMSE for larger values of $M$: it takes longer for a model trained on larger $M$ to stop moving towards Ridge and turn around and move towards dMMSE. However, we note that the hypothesis of undertraining generated from this observation is unlikely to explain the task diversity threshold given that we train for **much longer** than the longest time it takes for models in group (2) to turn away from Ridge. Concretely, the green curve for $M=2^{10}$ in Fig 2A right turns around at $t^*$ ~ 100K steps. The red curve for $M=2^{11}$ increase $M$ by a factor of 2 and yet, continues to move towards Ridge even though it is trained for 2M steps, more than 20x the $t^*$ noted for $M=2^{10}$. Thus the curves at high task diversity whose distance to ridge monotonically decreases are unlikely to suddenly turn around and depart from ridge after even longer training. To bolster this picture quantitatively, we perform a scaling analysis of the time $t^*$ it takes for curves in group (2) i.e at small M to turn around and start to move further from ridge with more training. This analysis reveals that this time $t^*$ scales roughly as $\\sqrt{M}$ (rebuttal Fig. 2B). However there is a sudden phase transition around the task diversity threshold where the distance to ridge curves start decreasing monotonically and don’t turn around even when trained for times **much** longer than the $\\sqrt{M}$ turn around time extrapolated from below the task diversity threshold (sudden large jump in the points in rebuttal Fig. 2B). For example, in rebuttal Fig. 2B, at $M=2^{11}$ task diversity, the $\\sqrt{M}$ scaling would predict the a turn around time around $t^*=10^5$. However, the training curve *doesn't* turn around for atleast 2 million training steps instead becoming monotonically closer to ridge.  This sudden and dramatically large change in behavior of the training curves between $M=2^{10}$ and $M=2^{11}$ vividly demonstrates the fundamental nature of the task diversity threshold phase transition and how it is highly unlikely to be an artifact of not enough training time at large $M$; for example training even longer at $M=2^{11}$ is unlikely to yield turn around behavior, and if it did, it would violate the $\\sqrt{M}$ scaling law in rebuttal Fig. 2B.  Finally, as further evidence (rebuttal Fig. 1)  we not only increased the amount of training data by 2x as in most of our experiments, **but also increased it by 8x** compared to our base setting: this did not change the task diversity threshold. \n\nTaken together these results provide strong evidence that our discovered task diversity threshold is not an artifactual consequence of our evaluation setup but rather it is inherent to the learning dynamics of the transformer.    \n\nThe reviewer asked why there is a task diversity threshold phase transition. While we don’t like to speculate, given the reviewer’s request we provide our best current guess in the global rebuttal, motivated by the  picture of the learning dynamics revealed  in Rebuttal Fig. 2. Please see global rebuttal for more details.         \n\n**Question 3:** No. Fig. 2, upper row: evaluated on finite set of pretraining tasks, lower row: evaluated on new tasks from Gaussian dist."}, {"review_id": "Li6HNXrVnU", "reviewer": "Reviewer_nyxK", "summary": "The 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).", "questions": "1. 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", "limitations": "The authors could discuss the possibility of this transition might not be hapenning in the case of enough iteration and samples per task.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. 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", "weaknesses": "1. 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\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Undertraining (Weakness 1 + Paper Summary):** To clarify, there are two independent parameters in our experimental setting that control how we scale the amount of data the model is trained on. We can vary the number of pretraining tasks, $M$, which is the x-axis in our experiments. We can also vary the total number of examples seen during training by increasing the batch size or training duration. The number of examples per task is the total number of examples seen divided by $M$. So for any fixed $M$, increasing the number of iterations increases the number of instances per task at the same time.\n\nHowever, there is still a question of if the task diversity threshold is caused by undertraining, which is what we think the reviewer’s question is about. It is possible that for a model to converge to dMMSE, when we increase $M$ by a factor of 2, we also need to increase the total number of training iterations by a factor of 2, thus keeping the total number of samples per task constant for each $M$. We discuss this question in detail with additional experiments in the global rebuttal. To summarize, we do a version of the “maxing out experiment” the reviewer asks for with a smaller model (which was necessary to finish the experiment within the week). For this Small PT, we train on 8x the amount of data (with larger batch size AND more iterations). We find that the location of the task diversity threshold is in fact robust. Since we performed a similar experiment in our paper (in Figure 2, we 4x the data when we go from batch size 256 to 1024), we believe that similar results will hold for larger models. We also provide additional evidence by looking at how the early stopping time for Ridge (the training step at which the model is closest to Ridge) scales with $M$; for details, please refer to the global rebuttal. This provides further evidence that there is a real change in the training dynamics of the models near the task diversity threshold and beyond the threshold, models really do prefer the Ridge solution no matter how long one trains them for.\n\n**Possible explanation for the transition (Weakness 2):** While we don’t like to speculate, given the reviewer’s request we provide our best current guess which is motivated by the picture of learning dynamics revealed in rebuttal Fig. 2. To summarize, we think that as the task diversity increases, there is a transition in the geometry of the loss landscape that makes the dMMSE solution too hard to find and the model instead picks the Ridge solution even though it has higher error. Please see the global rebuttal for a detailed discussion.\n\n**Exploring more hyperparameters (Weakness 4):** We further investigate the effect of capacity on this transition using our Small PT model (see global rebuttal). We hypothesize that a smaller model will be in a more capacity-constrained regime, so that ablations of model depth and embedding dimension will have a more pronounced effect on the task diversity threshold. If we start from Small PT and either increase the number of layers (rebuttal Figure 1B), or increase embedding size (rebuttal Figure 1C), the task diversity threshold goes up. This adds results for the effect of layers and presents evidence that model capacity can be a key feature of the task diversity transition at small model scales.\n\n**Laplace prior (Weakness 3 and Question 1):** We investigate if the task diversity threshold is also present when tasks are drawn from a different prior. We performed new experiments where pretraining tasks are sampled from a Laplace distribution, instead of a Gaussian (Rebuttal Fig. 3). In this setting we again observe a task diversity threshold in the model’s error where, beyond the threshold, more training leads to better performance on new, unseen Laplace tasks. We note that for $M>2^{12}$, the PT achieves a lower error on new Laplace tasks than ridge regression (~0.34). This suggests that the PT is primed to generalize on unseen tasks from the same underlying distribution that the pretraining tasks are sampled from. \n\n**Simplifying takeaways (Question 2 and 3):** We note training at very large task dimensions requires larger models to perform well and also exhibits many loss spikes which makes training unstable and very expensive. So given the number of training runs we need to conduct, experiments sweeping past $D=32$ become infeasible. For our task dimension experiments (Figure 5, main submission), we speculate that the scaling is linear, but don’t want to overclaim given that we’re sweeping task dimension over a factor of 4. However, we agree that the messages in our figures could be made clearer by presenting scaling laws for settings with sufficient datapoints. We will simplify the figures and corresponding takeaways in our paper; our latest architecture ablation experiments (rebuttal Figure 1B and C) already incorporate this suggestion."}, {"review_id": "YSkWIBEZfv", "reviewer": "Reviewer_NkJG", "summary": "The 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.   ", "questions": "Unfortunately, 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.  ", "limitations": "Limitations are obvious and mentioned in the discussion. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The 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\n", "weaknesses": "The 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments which we address below:\n\n**Novelty and significance of contribution (Weakness):** As the reviewer notes, it has been shown that if the pretraining distribution has infinite task diversity (each pretraining task is sampled i.i.d. from a Gaussian distribution), PTs implement Ridge regression. It is also perhaps unsurprising that if the pretraining distribution has a small number of tasks relative to the model capacity, the PT will memorize the tasks and implement dMMSE. However, we think it isn’t obvious that PTs would fall back to _Ridge regression_ if they do not have the capacity to _exactly implement dMMSE_. Specifically, consider a setting after the task diversity threshold, eg. at $2^{16}$ tasks in paper Fig. 2, where the _pretraining loss_ of the dMMSE estimator is significantly lower than that of Ridge. One natural hypothesis is that, since the PT is optimized to minimize the pretraining loss, even if it doesn’t have the capacity to exactly implement dMMSE, it will achieve a solution at an intermediate value of loss between dMMSE and Ridge. Under this hypothesis, if trained for longer, the PT would reduce its pretraining loss further, thereby becoming _less like_ Ridge ($\\Delta_\\text{PT,Ridge}$ would increase), until it converged to the closest solution to dMMSE that its capacity allows. However, we find that this is not the case: above the threshold the PT seems to “prefer” the Ridge solution. In our paper, and in the additional experiments in the rebuttal, we aim to establish that this “preference” is a robust feature and not just an artifact of undertraining (see the global rebuttal). To our knowledge, such a result has not been shown for transformers and we are not aware of results from classical statistical learning theory that would imply this without _assuming_ a biased function class. This finding is also significant as it shows an exciting example where, even when trained on finite task diversity (as one would expect to in practical settings), PTs show a bias for learning the underlying generative model of the tasks instead of just minimizing the pretraining loss.\n\nIn particular, we encourage the reviewer to see rebuttal Fig. 2 and the global rebuttal which reveals a remarkable and sharp phase transition in learning dynamics. To summarize: at small task diversity $M$, pre-training learning curves initially go towards ridge but then turn away and go to dMMSE. However, as $M$ increases, around the task diversity threshold, *suddenly*, learning curves escape dMMSE and move towards ridge *no matter how long you train despite the fact that the pre-training loss global minimum is dMMSE*. Rebuttal Fig 2B shows that this transition would not be predicted by a simple scaling analysis of early stopping time. We are unaware of classic results in statistical learning theory that can explain all these facts. We provide a conjectured loss landscape explanation in our global rebuttal. \n\n**Questions**:\n\n**Architecture modifications:** We further investigate architecture modifications and their effect on the transition in our Small PT. We hypothesize that a smaller model will be in a more capacity-constrained regime, so that the effects of architectural parameters on the task diversity threshold will be larger. Switching from Base PT to Small PT, which has embedding size 64, 4 layers, and 2 heads, leads to a task diversity threshold 8x lower. Furthermore, starting from Small PT, if we increase the number of layers (cf. rebuttal Fig. 1B), or embedding size (cf. rebuttal Fig. 1C), the task diversity threshold goes up. This shows that restricting the models class does in-fact reduce the task diversity threshold. Unfortunately, this also reduces the overall performance of the models, revealing a tradeoff between model performance and required diversity of dataset.\n\n**LSA:** The softmax might in fact play an important role in allowing the model to reach dMMSE-like solutions. We experimented with linear self-attention (LSA) and found training to be unstable, even after some hyperparameter tuning. A key difference that might cause this is that other works considering LSA only predict the final test example while we predict at each context length. Given that this isn't a simple modification, we think that investigating LSA is beyond the scope of this paper but an excellent opportunity for future work.\n\n**Open ended**:\n\n**Other task diversities.** We performed experiments where pretraining tasks are sampled from a Laplace prior, instead of a Gaussian (see Figure 3). In this setting we again find a task diversity threshold. Crucially, for $M>2^{12}$, the PT achieves a lower test error on new tasks drawn from the Laplace prior than does ridge regression (~0.34). This suggests that the PT is primed to generalize to unseen tasks from the underlying distribution that the pretraining tasks are sampled from. Thus, it might be tricky to engineer a different prior to sample tasks from to enable faster generalization. But this is a very promising direction for future work where the tasks come from hierarchical or otherwise structured priors that we might be able to exploit.\n\n**Lazy vs. feature learning:** It would indeed be interesting to use our setup as a tool to not just study implicit bias in this problem but also for getting a better understanding of lazy versus feature learning.\n\n**Evidence for copying and GD behavior in attention layers:** We have not had the opportunity this week to plot early attention layers. It’s certainly plausible that there will be noticeable differences for dMMSE vs. Ridge-like models, especially if, as the reviewer suggests, the PT is implementing a version of copy + GD in the post-transition regime. As alluded to in the LSA discussion, however, we’re not sure the solution would appear structurally as simple in our setting of large models with MLP’s and softmax attention, as well as predicting at every context length."}, {"review_id": "RLbkjllDyH", "reviewer": "Reviewer_Uwh7", "summary": "The 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.", "questions": "1. (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", "limitations": "1. 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.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses:**\n- We agree it’s important to add a discussion positioning our work with respect to the few-shot and meta-learning settings, particularly considering the role of task diversity. Thank you for the references; we will add them to our paper and discuss them and others. \n- Please refer to the main rebuttal response for a discussion on underfitting (also as mentioned in Question 1).\n\n**Effect of model capacity (Question 1):** We further investigate the role of capacity in our Small PT. We hypothesize that a smaller model will be in a more capacity-constrained regime, so that ablations of key parameters affecting capacity will have a more pronounced effect on the task diversity threshold. Our Small PT has embedding size 64, 4 layers, and 2 heads. If we start from Small PT and either increase the number of layers (Figure 1B), or increase embedding size (Figure 1C), the task diversity threshold goes up. This presents evidence that model capacity can be a key feature of this transition at small model scales.\n\nWith “A PT that minimizes the pretraining loss $L^{T_\\text{Pretrain}}$ will behave like this estimator,” we mean that a model that is able to attain (close to) the lowest possible $L^{T_\\text{Pretrain}}$ over all functions will behave like the dMMSE estimator. For large enough $M$, a model _trained to minimize_ $L^{T_\\text{Pretrain}}$ in fact does not find this minimum and so does not behave like dMMSE; we will word this more clearly in our paper.\n\n**Defining fundamentally new tasks (Question 2):** We first clarify that new test tasks used to compute $L^{T_\\text{True}}$ are sampled from a different distribution than pretraining ones: to generate a pretraining sequence, we sample a task from the _uniform_ distribution over pre-selected pretraining tasks $\\{w_1,\\ldots, w_M\\}$, whereas when we calculate $L^{T_\\text{True}}$, we average over the underlying task distribution $\\mathcal{N}(0, I)$.\n\nIn this setup we can make the notion of “fundamentally new” precise because we have a well-defined metric on task space: tasks are regression vectors and distance is simply cosine distance in this space. We show that distance of a new in-context task from a pre-training task strongly determines performance below the task diversity threshold but not above it. This is shown in Figure 4 (main submission) when we interpolate between pre-training tasks and measure in-context learning performance.\n\n**Distribution shift (Limitation 1):** It is the case that the pretraining tasks $\\{w_1,\\ldots,w_M\\}$ are sampled from the same _underlying_ Gaussian distribution as test tasks. It is likely that this property does not hold in language modeling settings where the pretraining tasks and test tasks themselves come from different distributions. We will discuss this further as a limitation in our paper. However, we think that understanding the simpler setting where the distributions are the same is an important first step. Searching for settings in which there is a shift in the underlying distribution and near-optimal generalization persists is an exciting direction for future work.\n\n**Q3: Connection to supervised learning and IID generalization:**\nIndeed, this problem can be framed as an instance of supervised learning where at each context, $k$, the datapoint is $S_k$ and the target is $y_k$. Standard IID generalization in our setting, however, corresponds to strong performance on $L^{T_\\text{Pretrain}}$, for which test sequences are distributed identically to pretraining ones, with tasks sampled from a finite pool of $M$ regression vectors. In fact, dMMSE is the optimal model for _IID_ generalization. Conversely, good performance on $L^{T_\\text{True}}$ means that the PT is generalizing _out of distribution_ in that it is able to perform well on _new_ tasks drawn from $\\mathcal{N}(0, 1)$ it was not trained on. The key is that, even though the set of pretraining tasks were sampled from the Gaussian, once those are fixed and we start pretraining, the distribution of $(S_k, y_k)$ is different from the distribution of $(S_k, y_k)$ used to calculate $L^{T_\\text{True}}$. What makes this behavior emergent is precisely that even though the model is trained to minimize $L^{T_\\text{Pretrain}}$, it ends up performing well on $L^{T_\\text{True}}$, and at the expense of not minimizing $L^{T_\\text{Pretrain}}$ (see Figure 2, top left in main text). Moreover training with *more* data per task makes the model more like Ridge at high task diversity and more like dMMSE at low task diversity.  This asymptotic approach to two *different* task-diversity dependent estimators given more training data is a distinct phenomenon that does not show up in classic learning theory (see our new results in Rebuttal Fig. 2AB and global rebuttal). "}, {"review_id": "8uJB3WxIN8", "reviewer": "Reviewer_CEjq", "summary": "This 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.", "questions": "N/A", "limitations": "Yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive impression of our paper and hope they will champion it.\n\nWe agree with the reviewer’s comments that while this formulation is simple and the connection to language modeling is less clear,  understanding the fundamental role of pretraining data/task diversity in this simpler setting will provide important intuition for more complex problems. \n\nAlong these lines, we note intriguingly,  an interesting recent paper that attempts to develop a data diversity metric for language data: “Beyond Scale: the Diversity Coefficient as a Data Quality Metric Demonstrates LLMs are Pre-trained on Formally Diverse Data” https://arxiv.org/abs//2306.13840.  We became aware of this paper after our submission.  This paper is along the lines of what the reviewer is asking for: a connection between dataset diversity and performance within the context of language and LLMs.  It remains for further exploration the relationship between our results on task diversity and its impact on in-context learning performance on new regression tasks, and the new paper’s data diversity metric and its ability to predict in-context learning performance on new language tasks.  However, exciting future work will be to explore the role of data diversity and its impact on transformer learning in many domains ranging from the relatively theoretically tractable settings of regression in which we know how optimal solutions behave, to the relatively open ended domains of language, vision, and multimodal data, in which new measures of dataset diversity that are predictive of performance may need to be derived.  For example, our work may motivate exploring parametric variations in the data diversity coefficient of datasets used to pretrain transformer LLMs to see if their learning performance exhibits a phase transition in data diversity - this would allow us to find perhaps small datasets of minimal diversity just above the threshold required to solve many language tasks.  We will add these ideas to our discussion section of the camera ready version. \n\nFinally, since we work with multi-layer transformers and softmax self-attention, a direct theoretical analysis in this exact setting is challenging. However, we aimed to firmly and rigorously establish the existence of the task diversity threshold empirically and understand how it behaves as we change dimensionality of the problem (paper Fig 5), regularization (paper Fig 6, weight decay) and capacity (paper Fig 6, embedding and new experiments in rebuttal). We hope this will serve as a strong starting point for future theoretical studies."}]

Author rebuttal:
We thank the reviewers for their insightful comments. 
Summary of new results in the attached figures:
- We study a Small PT with 64 dim. embeddings, 4 layers, and 2 attn. heads; the smaller scale enabled us to answer questions within the week. In Fig 1A, we reproduce our key finding, the existence of a task diversity threshold.
- A Small PT, with a lower capacity, has a lower task diversity threshold than the Base PT in our paper. Since PTs at this scale are more sensitive to capacity constraints, we vary depth and embedding size on Small PTs showing that as capacity increases, so does the task diversity threshold. The results (Fig 1 B,C) are discussed in individual responses.
- Multiple reviewers asked if the task diversity threshold may be an artifact of undertraining. We train Small PT for much longer and address this question below; it is not the case.
- We investigate the effect of changing the task distribution from Gaussian to Laplace. In Fig 3, we show that a task diversity threshold also exists in this setting. We discuss this further in the response to reviewer nyxK who raised this question.

**Is the task diversity threshold an artifact of undertraining?** In rebuttal Fig. 2A, we see that for some pretraining task diversities, $M$, below the task diversity threshold, $M^*$, $\Delta_\text{PT,Ridge}$ decreases early in training until it reaches a minimum at time step $t^*$, and then increases as the PT converges to dMMSE ($M=2^8,2^9,2^{10}$ in Fig. 2A). Call $t^*$ the early stopping time for Ridge. A reviewer noted that a possible reason for the task diversity threshold could be that for $M > M^*$, $t^*$ is larger than the total training time, and so $\Delta_\text{PT,Ridge}$ appears to decrease with more training. Indeed, in an ideal experiment, we would keep the number of examples per task fixed as we increase $M$, thus training on more data for larger $M$. Unfortunately, this is infeasible for the range of $M$ we study. However, we argue that the task diversity threshold is a feature inherent to the problem and not merely an artifact of undertraining and present the following evidence:
1. For each $M$, we train a Small PT for 2M steps at batch size 512 (rebuttal Fig. 1A and Fig. 2A right panel). Note, this is 8x as many examples per task as our base setting (500K steps at batch size 256). Despite this, the location of the task diversity threshold does not change. If the task diversity threshold were to disappear in the ideal experiment described above, then in the experiment we presented, we would expect the threshold to increase at least by a factor of 2 when we increase training data from 2x to 8x examples per task. But this is not the case, suggesting that the location of the threshold is robust in this setting.
2. Of course, this does not eliminate the possibility that $t^*>$ 2M: with finite training time experiments, we cannot prove that models after the task diversity threshold will not eventually converge to dMMSE. However, in rebuttal Fig. 2B, we study how the early stopping time for Ridge, $t^*$, scales with $M$. For most $M<M^*$, $t^*$ obeys a simple scaling behavior $t^* \propto M^\alpha$ with $\alpha \approx 0.47$. However, for $M>2^{10}$, the distance to Ridge decreases monotonically through training and $t^*=$2M. While it is possible that, at some $t>$ 2M steps, the model’s distance to Ridge will start increasing again, the $t^*$ would be much larger than what is predicted by the scaling relationship. The break in the scaling behavior near the task diversity threshold suggests that our findings are not just caused by undertraining. There is a difference in the training dynamics and the dMMSE solution, if not unreachable, at least takes an uncharacteristically long time to find.
3. In Fig 2A, we visualize the training curves for the models at batch size 512 for both 500K and 2M steps for $M$ before and after the threshold as well as the training curve for the model trained with infinite task diversity. For both the shorter and longer training durations, models trained with the same task diversity have similar qualitative behavior (whether distance to Ridge decreases then increases or monotonically decreases). Additionally, the learning curves of the models beyond the task diversity threshold are very similar to the learning curves for models trained on infinite pretraining task diversities (Fig 2A) and they achieve similar final accuracies (dashed lines vs markers in Fig 1A), suggesting that these models really are approaching the ridge solution, and will not depart from it even if trained longer.

**A possible picture/explanation for the task diversity threshold.** Multiple reviewers have asked why there is a task diversity threshold phase transition. While we don’t like to speculate, at the reviewers’ request, we provide our best current guess given the new rebuttal experiments. Indeed the result in Fig. 2 suggests a natural loss landscape explanation for the task diversity threshold phase transition. The pre-training loss landscape always has a local minimum corresponding to dMMSE. We conjecture: at low task diversity ($M < 2^{11}$), this local minimum is wide enough to eventually attract the training trajectory. We further conjecture: as task diversity increases, the local minimum corresponding to dMMSE becomes sharper/its basin of attraction becomes smaller until it can no longer attract training trajectories. After the task diversity threshold, models do not find the sharp dMMSE solution no matter how long you train, but instead, go to a higher training error but wider minimum that corresponds to Ridge. While this is speculation, it does suggest many interesting loss landscape analyses that would be exciting follow up works to our work. In this sense, we hope our discovery and substantiation of a task diversity threshold phase transition constitutes an important and impactful contribution to the NeurIPS community and inspires follow up work.

Meta-review: The submission contributes hypothesis-driven experiments to a highly-discussed subarea (in-context learning / few-shot generalization in transformers). Though several reviewers noted the gap between the setting examined in the submission and the broader settings in which transformers are employed, all reviewers found the submission thought-provoking with the potential to spur more conclusive investigations into why the factor identified in the work (data diversity in pre-training) matters.
Decision: Accept (poster)