SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Finite size scaling of training data suggests an algorithmic phase transition as task-diversity increases. The experiments in Fig. 2 (left column) suggest that, when tested on both task distributions  $\mathcal{T}_{\mathrm{Pretrain}}$  and  $\mathcal{T}_{\mathrm{True}}$ , the ICL algorithm implemented by a PT exhibits a smooth crossover in performance from dMMSE to Ridge. We next examine how this transition changes as we increase the number of sequences per task seen over pretraining, at fixed task diversity. One might reasonably expect that, if the transformer sees more sequences per latent task in  $\mathcal{T}_{\mathrm{Pretrain}}$ , both its predictions and performance should become more similar to those of dMMSE, and less similar to those of Ridge, at all values of task diversity. Strikingly, this natural expectation is violated in a manner that facilitates ICL on  $\mathcal{T}_{\mathrm{True}}$ .
At each number of tasks, we increase the number of sequences per task by increasing batch size from 256 to 512 to 1024, while leaving the number of training steps fixed at  $500\mathrm{K}$ . We observe that  $\Delta_{\mathrm{PT,dMMSE}}^{\mathcal{T}_{\mathrm{Pretrain}}}$ , which quantifies how different the PT and dMMSE estimator's predictions are when testing on tasks drawn from  $\mathcal{T}_{\mathrm{Pretrain}}$ , does in fact decrease for  $M \leq 2^{10}$  (Fig. 2 top center) as we train on more sequences per task. Moreover, for each  $M \in \{2^{10}, \dots, 2^{14}\}$  the PT's predictions also become less similar to those of Ridge, both on tasks from  $\mathcal{T}_{\mathrm{Pretrain}}$  (Fig. 2, top right) and  $\mathcal{T}_{\mathrm{True}}$  (Fig. 2, bottom right). Crucially, this movement in behavior of the PT towards dMMSE and away from Ridge, at least on tasks drawn from  $\mathcal{T}_{\mathrm{Pretrain}}$ , holds only up to a threshold number of tasks between  $2^{14}$  and  $2^{15}$ . Beyond this threshold, pretraining on more sequences per task at a fixed task diversity actually makes the PT more like Ridge, in that both  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{Pretrain}}}$  and  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  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  $\mathcal{T}_{\mathrm{True}}$  by matching Ridge performance, but also the PT gets better at doing so if trained on more sequences per task, despite the limited set of tasks experienced in pretraining. Thus, in contrast to the natural expectation stated above, more sequences per task does not promote overspecialization of the PT to the  $\mathcal{T}_{\mathrm{Pretrain}}$  at task diversities beyond the threshold.
Finally, the motion of the ICL algorithm implemented by PT towards (away) from Ridge above (below) a task diversity threshold (Fig. 2, right top and bottom) indicates that as one increases the number of sequences per task at fixed task diversity, the smooth cross over in performance of the PT between dMMSE and Ridge, shown in Fig. 2, left top and bottom, will become sharper and sharper in task diversity, ultimately exhibiting a sharp phase transition in the limit of infinite number of sequences per task. Remarkably, this phase transition in the ICL algorithm implemented by the PT appears at a moderate task diversity threshold below  $2^{15}$  pretraining tasks; even though dMMSE significantly underperforms relative to Ridge on  $\mathcal{T}_{\mathrm{True}}$  at this task diversity, the PT nevertheless remains unimpaired by this limited task diversity and can optimally solve new tasks.
Increased training time at fixed batch size further supports an algorithmic phase transition. To confirm the above results, we also increase the number of sequences per task, at each task diversity, by increasing the number of training steps  $N$  from 500K to 1M while keeping batch size fixed at 256. We observe that doubling  $N$  (change from pale blue to red in Fig. 3) and doubling  $B$  (change from pale
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^{5}$  tasks, well below the task diversity threshold, the PT performance matches that of the dMMSE estimator along interpolating paths, but under-performs Ridge on new tasks at the center. Middle: At  $2^{10}$  tasks, the PT outperforms dMMSE on new tasks at the center of the interpolation path, but is not yet as good as Ridge on new tasks. Right: At  $M = 2^{15}$  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 a finite but large number of unique tasks, the PT, unlike the Bayes optimal estimator for  $\mathcal{T}_{\mathrm{Pretrain}}$ , can learn new tasks in-context.
blue to red in Fig. 2) have very similar effects on  $\Delta_{\mathrm{PT,dMMSE}}^{\mathcal{T}}$  and  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}}$ , for both  $\mathcal{T} = \mathcal{T}_{\mathrm{True}}$  and  $\mathcal{T} = \mathcal{T}_{\mathrm{Pretrain}}$ . More importantly, the task diversity threshold, which we determined as the cross-over point in  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  between batch sizes 256, 512, and 1024 at 500K training steps (Fig. 2 bottom right) happens at the same number of tasks as the crossover point between 500K and 1M steps at batch size 256 (Fig. 3, right). Given that our two approaches for training the baseline transformer on more data yield the same task diversity threshold, and that doubling batch size leads to significantly faster training times than doubling number of steps, from here onward we consider the task diversity threshold to be cross-over point in  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  between batch sizes 256 and 512 when training for 500K steps. See Appendix D for more ablations of batch size and training steps that provide further evidence for how the number of sequences seen by the transformer is the key factor determining the similarity of its predictions to those of dMMSE and Ridge at each number of tasks.
Learning dynamics and a break in the scaling of early stopping time further supports an algorithmic phase transition. To probe if the observed transition is merely an effect of under-fitting, we study the learning dynamics of small PTs in the very large number of steps regime. First, in Appendix E, we verify that the small PT also demonstrates an algorithmic phase transition but at lower task diversity threshold between  $2^{11}$  and  $2^{12}$  pretraining tasks. In Fig. 4 left, we visualize the learning curves  $(\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  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^{*}$ ,  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  decreases early in training until it reaches a minimum at time step  $t^{*}$ , and then increases as the PT approaches dMMSE. We define  $t^{*}$  as the early stopping time for Ridge. For  $M > M^{*}$ ,  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  decreases throughout training. To evaluate if, in the latter case, models are undertrained and  $t^{*}$  is larger than the total training time, we extend training to 2M steps at batch size 512 ( $4\times$  the training time, see Appendix B). Fig. 4 center, shows these learning curves along with that of the model trained with infinite task diversity; even in this long training regime, the task diversity threshold does not change. For both short and long training durations, models trained with the same  $M$  have similar qualitative behavior (whether distance to Ridge decreases then increases or monotonically decreases). Additionally, learning curves of the models with  $M > M^{*}$  are very similar to the learning curves for models trained on infinite pretraining task diversities and they achieve similar final accuracy (dahed lines vs markers in Fig. 10), suggesting that these models are approaching the Ridge solution.
In Fig. 4 right, we study how  $t^*$ , scales with  $M$ . For most  $M < M^*$ ,  $t^*$  obeys a simple scaling behavior  $t^* \propto M^\alpha$ ,  $\alpha \approx 0.47$ . However, for  $M > 2^{10}$ , the distance to Ridge decreases monotonically through training and  $t^* = 2M$  steps. Despite the caveat that our experiments are necessarily in the large but finite training step regime with a decayed learning rate schedule, this stark break in the scaling behavior of  $t^*$  near the task diversity threshold suggests that the observed transition is not just caused by under-fitting but an underlying difference in the learning dynamics.
Figure 6: 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  $D$  (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  $2^{20}$  tasks (right), the PT's predictions are similar to those of Ridge at all  $D$  (orange  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$ ), whereas dMMSE grows progressively less similar to Ridge (blue  $\Delta_{\mathrm{Ridge,dMMSE}}^{\mathcal{T}_{\mathrm{True}}}$ ).
The transition along interpolating paths. To obtain an additional description of the algorithmic transition in the PT from dMMSE to Ridge, we compute the ICL performance of the PT, and compare it to both dMMSE and Ridge, on a one parameter family of new tasks  $\mathbf{w}_{\alpha}$  that interpolate between pairs of seen tasks  $\mathbf{w}_i$  and  $\mathbf{w}_j$  in the support of  $\mathcal{T}_{\mathrm{Pretrain}}$ . The interpolation path is given by
$$
\mathbf {w} _ {\alpha} = \frac {1}{2} \left(\| \mathbf {w} _ {i} \| _ {2} + \| \mathbf {w} _ {j} \| _ {2}\right) \frac {\alpha \mathbf {w} _ {i} + (1 - \alpha) \mathbf {w} _ {j}}{\| \alpha \mathbf {w} _ {i} + (1 - \alpha) \mathbf {w} _ {j} \| _ {2}} \quad \text {f o r} \alpha \in [ 0, 1 ]. \tag {5}
$$
Here we fix the norm of the interpolated vector  $\mathbf{w}_{\alpha}$  to the average of the two endpoint norms to avoid  $\| \mathbf{w}_{\alpha}\|$  taking on very small values for  $\alpha \sim \frac{1}{2}$ . Fig. 5 shows the results of this analysis for  $2^{5}$  (left, low task diversity regime),  $2^{10}$  (center, below task diversity threshold), and  $2^{15}$  (right, just above the task diversity threshold) tasks. At each value of  $\alpha$ , MSE is averaged over a large number of task pairs  $(\mathbf{w}_i,\mathbf{w}_j)$ . Examination of the average performance at the center of many interpolation paths, corresponding to fundamentally new tasks far from tasks seen during pretraining, clearly reveals a transition in PT performance from dMMSE to Ridge, where new tasks can only be optimally learned above, but not below, the task diversity threshold. In contrast, unlike the PT, dMMSE cannot solve new tasks at any task diversity in the range considered.
The PT outperforms a smoothed dMMSE model. We have seen that at an intermediate task diversity the PT significantly outperforms dMMSE on new tasks in  $\mathcal{T}_{\mathrm{True}}$ . It is clear why dMMSE performs poorly on new tasks in  $\mathcal{T}_{\mathrm{True}}$  at low task diversity: its prior over tasks concentrates on  $M$  unique tasks in  $\mathcal{T}_{\mathrm{Pretrain}}$ , while the prior over tasks in  $\mathcal{T}_{\mathrm{True}}$  is Gaussian. A natural conjecture is that the PT cannot memorize all  $M$  tasks in  $\mathcal{T}_{\mathrm{Pretrain}}$  for large enough  $M$ . Therefore we also compare PT performance to a smoothed dMMSE estimator in which the discrete point prior over  $M$  tasks seen in pretraining is replaced with a mixture of  $M$  isotropic Gaussians with the same centers but with variance chosen to optimize performance on  $\mathcal{T}_{\mathrm{True}}$  (see Appendix G for details). This smoothed dMMSE outperforms dMMSE as it has a prior over tasks closer to the Gaussian  $\mathcal{T}_{\mathrm{True}}$ . But remarkably, the PT still outperforms the smoothed dMMSE even with optimal smoothing (Fig. 12). This indicates the PT, even at moderate task diversity, implements a more sophisticated algorithm than a simple smoothed dMMSE arising from the PT's inability to resolve the  $M$  pretraining tasks to high precision.
3.2 The PT exhibits superior scaling of task diversity threshold with dimension than dMMSE.
We next explore the dependence of the task diversity threshold on the regression problem dimension  $D$ . We vary  $D \in \{8, 16, 32\}$  while simultaneously scaling up maximal context length as  $K = 2D$ , and increasing observation noise  $\sigma^2$  to match the SNR to that of  $D = 8$  and  $\sigma^2 = 0.25$ . We also train a larger model with 12 layers, 256-dimensional embeddings, and 4 attention heads that is sufficiently expressive to match Ridge performance at  $D = 32$ . Fig. 6, first 3 panels reveal that the task diversity threshold of the PT increases moderately (approximately linearly) with task dimension (i.e. roughly  $2^{14}, 2^{15}$ , and  $2^{16}$  at  $D = 8, 16,$  and 32 respectively). This linear scaling is remarkable considering the volume of all possible tasks scales exponentially with dimension due to the concentration of the Gaussian  $\mathcal{T}_{\mathrm{True}}$  to a sphere for large  $D$ . Thus we expect dMMSE performance to scale much more poorly with dimension  $D$  since the finite number of tasks in  $\mathcal{T}_{\mathrm{Pretrain}}$  would need to cover a substantial portion of the sphere for dMMSE to approach Ridge. To test this hypothesis, for  $M = 2^{20}$ , which is
Figure 7: Explicit regularization and model capacity affect the task diversity threshold. Increasing 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  $\Delta_{\mathrm{dMMSE,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  vs  $M$ .
the largest task diversity we consider, we explore how the similarity of PT and dMMSE predictions to Ridge on new tasks scales with  $D$  (Fig. 6, right panel). We see that  $\Delta_{\mathrm{dMMSE,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  grows significantly as we increase  $D$ , while remarkably  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}_{\mathrm{True}}}$  is largely dimension independent. Overall this indicates that the scaling of PT error with dimension is vastly superior than that of dMMSE; PT remains near optimal and close to Ridge at all  $D$  for  $M = 2^{20}$ , while dMMSE departs from Ridge as  $D$  increases.
3.3 Effect of Regularization and model capacity on the task diversity threshold.
We study the dependence of the task diversity threshold on various hyperparameters. First, adding explicit regularization in the form of weight decay (see Appendix B for details), and increasing its value over three orders of magnitude, consistently lowers the threshold task diversity (Fig. 7, left). Note however, the lower task diversity threshold also comes with worse performance (Figure 13, top). This suggests that various forms of implicit regularization could help drive the algorithmic transition in the PT without weight decay. We also explore the effect of model capacity on the task diversity threshold by either increasing the embedding dimension of both small and base PTs or increasing the depth of small PTs. Fig. 7 center shows that increasing embedding dimension over a reasonable range does not affect the task diversity threshold of base PT. However, for small PT, increasing either the embedding dimension (Fig. 7 center) or depth (Fig. 7 right) increases the task diversity threshold. Base PT has a much larger capacity than small PT and also has a larger threshold; we hypothesize that small PT is still in a regime where the threshold is sensitive to capacity while base PT is not. Together, these results suggest that model capacity plays an important role in the emergence of in-context learning: increasing capacity (up to a point) leads to an increase in the task-diversity threshold.
The Bayesian framework for ICL introduced by Xie et al. [3], which motivates our work, hypothesizes that PTs "locate" concepts learned during pretraining to solve ICL tasks. A series of empirical work in language models [18-20] use this framework to select better in-context examples while Min et al. [21] use it to study the robustness of latent task inference. Our work builds on this framework in the linear regression setting and validates it at low task diversities. However, we find a regime—large but finite number of pretraining tasks—in which the ability to learn new tasks in-context is an emergent phenomenon that cannot be fully explained by Bayesian inference.
Prior work [12-14] has also shown that transformers can do linear regression in-context. However, they pretrain with unlimited task diversity, sampling a completely new regression vector for each sequence. In contrast, our work considers pretraining datasets with limited task diversity where ICL on new tasks emerges even though the pretraining loss does not explicitly encode it. Another line of work hypothesizes that ICL performs gradient descent in the activations of the forward pass, providing explicit constructions for the weights of the PT to implement this for linear regression [13, 14] or exploring this hypothesis in language models [22]. However more experiments are required to test the hypothesis that trained transformers actually match proposed constructions. Instead of studying the explicit mechanism by which in-context learning is implemented, our work focuses on the impact of the pretraining task diversity. Similar questions pertaining to the role of task diversification have been explored in the meta-learning literature [23-25].
Kirsch et al. [26] also show the emergence of in-context learning with pretraining task diversity on a toy classification task. By studying this question in the controlled setting of linear regression, we can compare to the optimal estimators on  $\mathcal{T}_{\mathrm{Pretrain}}$  and  $\mathcal{T}_{\mathrm{True}}$ . This allows us to establish that ICL at finite task diversity emerges because the PT departs from the optimal estimator on the pretraining task distribution, and is not just a consequence of the pretraining task distribution becoming similar to the ideal task distribution. Among other important perspectives on ICL, Chan et al. [27] identify, in a toy setting, several properties of the training distribution—burstiness and occurrence of rare classes—that are necessary for the emergence of ICL. Wei et al. [28] study how ICL in large language models is affected by semantic priors and input-label mappings, focusing on differences across model scale. Olsson et al. [29] study inductions heads—circuits responsible for completing patterns by copying tokens—as a mechanism for implementing ICL.
Overall, we have extensively explored the impact of pretraining task diversity on 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  $\mathcal{T}_{\mathrm{Pretrain}}$  (i.e. dMMSE), or for the diverse ideal task distribution  $\mathcal{T}_{\mathrm{True}}$  (i.e. Ridge). These comparisons reveal an algorithmic phase transition in the PT from the former to the latter at an intermediate task diversity threshold; beyond this threshold, the PT solves fundamentally new tasks not seen during pretraining. Strikingly, this task diversity threshold scales moderately with task dimension, over the range of dimensions considered, despite the exponential growth in the volume of all possible tasks with dimension. Indeed this PT scaling vastly outperforms that of dMMSE. Overall, these results indicate that ICL of new tasks by PTs is an emergent phenomenon that cannot be explained by Bayesian inference on limited diversity pretraining task distributions. Moreover, our experiments suggest some form of implicit regularization in PTs allows them to break free of the pretraining task distribution to solve new tasks, given a moderate pretraining task diversity.
Remarkably, beyond the task diversity threshold, PTs learn the optimal estimator for the underlying generative model for pretraining tasks; this is the case for both Gaussian and Laplace priors over tasks (see Figure 14 for experiments with Laplace prior). This is true even though solutions with lower training loss exist; indeed when trained on more data at fixed diversity, PTs behave more like Ridge at the expense of higher training loss. Our experiments in Fig. 4 suggest that this algorithmic transition is due an underlying change in learning dynamics. We explore this hypothesis by probing the linear mode connectivity of the loss landscape [30, 31]. In Fig. 11 we find that PTs trained with large  $M$  inhabit the same loss basin as PTs trained with  $M = \infty$ : the training loss barrier between PTs trained with  $M \gtrsim 2^{13}$  and PTs with  $M = \infty$  is similar to two PTs trained with  $M = \infty$ . In contrast, there are large loss barriers between PTs trained with  $M < 2^{13}$  and  $M = \infty$ . Additionally, PTs trained with  $M = \infty$  are closer in weight space to PTs trained with large  $M$  than those trained with small  $M$  (see Appendix F). Overall, these experiments provide further evidence that PTs trained with task diversities beyond the threshold find solutions similar to the optimal model for  $\mathcal{T}_{\mathrm{True}}$ ; we leave further exploration of these loss landscapes to future work.
An intriguing question is how these observations carry over to language. A key mystery about the efficacy of ICL in language tasks lies in how different the tasks learned in-context are from the pretraining distribution of large language corpora. It is also less clear how to categorize the contents of such corpora according to tasks and measure their resulting task diversity. Regardless, our observation in linear regression that a moderate threshold in pretraining task diversity can enable PTs to solve new tasks may imply that many language tasks that are quite different from the statistics of large language corpora can still nevertheless be solved in-context.
Our results also suggest that the scale of data alone does not lead to good ICL performance. In fact, below the task diversity threshold, increasing the size of the pretraining dataset without increasing task diversity hurts ICL performance. It is necessary to increase both the diversity and size of the dataset for ICL to emerge. Thus to improve ICL in language settings, our work motivates future studies into uncovering the relevant notion of tasks in language modeling and approaches to increase task diversity in language corpora. More generally, our empirical analysis of the impact of pretraining task diversity on ICL motivates further theoretical studies. Such studies will be key to understanding the mystery of why simple next token prediction during pretraining can lead to in-context learning of so many apparently different tasks.
The authors would like to thank the Google TPU Research Cloud (TRC) whose generous support made this project possible. SG thanks an NSF CAREER award and NTT Research for support.