SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[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 as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023. URL https://openreview.net/forum?id=fzbHRjAd8U.
[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 representations, 2022.
[26] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers, 2022.
[27] Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya Singh, Pierre H. Richemond, Jay McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers, 2022.
[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, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.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 a winning ticket's mask? In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=xSsW2Am-ukZ.
A Bayesian estimators
A.1 Bayesian MMSE estimator
The optimal estimator for the  $k$ th prediction,  $\hat{y}_k^{\mathcal{T}}$ , that minimizes the  $k$ th term in the sum in  $L^{\mathcal{T}}$  (in Eq. (1)) is the Bayesian estimator with  $\mathcal{T}$  as the prior. This is given by the posterior mean of  $y_k$  conditioned on the context:  $\hat{y}_k^{\mathcal{T}} = \mathbb{E}_{\mathcal{T}, \varepsilon_k} [y_k \mid S_k]$ , where the expectation is over the task distribution,  $\mathcal{T}$ , and the noise,  $\varepsilon_k$ . To derive this, notice that, using the law of total expectation, the  $k$ th term in the loss  $L^{\mathcal{T}}$  is:
$$
\underset { \begin{array}{c} \mathbf {w} \sim \mathcal {T} \\ \mathbf {x} _ {1}, \ldots , \mathbf {x} _ {k} \sim \mathcal {N} (\mathbf {0}, \mathbf {I} _ {D}) \\ \varepsilon_ {1}, \ldots , \varepsilon_ {k} \sim \mathcal {N} (0, \sigma^ {2}) \end{array} } {\mathbb {E}} \left(f _ {\theta} (S _ {k}) - y _ {k}\right) ^ {2} = \underset {S _ {k}} {\mathbb {E}} \underset {\varepsilon_ {k} \sim \mathcal {N} (0, \sigma^ {2})} {\mathbb {E}} \left[ \left(f _ {\theta} (S _ {k}) - y _ {k}\right) ^ {2} \Big | S _ {k} \right]
$$
which implies that the optimal estimator is  $\hat{y}_k^{\mathcal{T}} = \mathbb{E}_{\mathcal{T},\varepsilon_k}[y_k\mid S_k]$ . Adopting the notation  $\mathbf{X} = (\mathbf{x}_1^\top ,\dots,\mathbf{x}_{k - 1}^\top)\in \mathbb{R}^{(k - 1)\times D}$  and  $\mathbf{y} = (y_{1},\ldots ,y_{k - 1})$ , we can express the expectation explicitly:
$$
\begin{array}{l} \mathbb {E} \left[ y _ {k} \mid S _ {k} \right] = \int d \mathbf {w} d y _ {k} y _ {k} p (\mathbf {w}, y _ {k} \mid \mathbf {X}, \mathbf {y}, \mathbf {x} _ {k}) \\ = \int d \mathbf {w} d y _ {k} y _ {k} p \left(y _ {k} \mid \mathbf {x} _ {k}, \mathbf {w}\right) p \left(\mathbf {w} \mid \mathbf {X}, \mathbf {y}\right) \\ = \int d \mathbf {w} \mathbf {w} ^ {\intercal} \mathbf {x} _ {k} p (\mathbf {w} \mid \mathbf {X}, \mathbf {y}) \\ \equiv \hat {\mathbf {w}} _ {k} ^ {\intercal} \mathbf {x} _ {k} \\ \end{array}
$$
$$
\hat {\mathbf {w}} _ {k} = \mathbb {E} [ \mathbf {w} \mid S _ {k} ] = \frac {\int d \mathbf {w} p (\mathbf {w}) \mathbf {w} \prod_ {i = 1} ^ {k - 1} p \left(y _ {i} \mid \mathbf {x} _ {i} , \mathbf {w}\right)}{\int d \mathbf {w} p (\mathbf {w}) \prod_ {i = 1} ^ {k - 1} p \left(y _ {i} \mid \mathbf {x} _ {i} , \mathbf {w}\right)} \tag {6}
$$
A.2 dMMSE estimator
For task distribution  $\mathcal{T}_{\mathrm{Pretrain}} = \mathcal{U}\{\mathbf{w}^{(1)},\dots,\mathbf{w}^{(M)}\}$ , we can directly get the discrete minimum mean squared error (dMMSE) estimator by plugging the uniform discrete distribution into Eq. 6.
$$
\hat {\mathbf {w}} _ {k} ^ {\mathrm {d M M S E}} = \sum_ {i = 1} ^ {M} \frac {\exp \left(- \frac {1}{2 \sigma^ {2}} \sum_ {j = 1} ^ {k - 1} \left(y _ {j} - \mathbf {w} ^ {(i) \intercal} \mathbf {x} _ {j}\right) ^ {2}\right)}{\sum_ {l = 1} ^ {M} \exp \left(- \frac {1}{2 \sigma^ {2}} \sum_ {j = 1} ^ {k - 1} \left(y _ {j} - \mathbf {w} ^ {(l) \intercal} \mathbf {x} _ {j}\right) ^ {2}\right)} \mathbf {w} ^ {(i)}. \tag {7}
$$
A.3 Ridge estimator
For task distribution  $\mathcal{T}_{\mathrm{True}} = \mathcal{N}(\mathbf{0},\mathbf{I}_D)$ , we can directly get the Ridge regression estimator from Eq. 6:
$$
\begin{array}{l} \hat {\mathbf {w}} _ {k} ^ {\text {R i d g e}} = \frac {\int \mathbf {w} \exp \left[ - \frac {1}{\sigma^ {2}} (\mathbf {X} \mathbf {w} - \mathbf {y}) ^ {T} (\mathbf {X} \mathbf {w} - \mathbf {y}) - \mathbf {w} ^ {T} \mathbf {w} \right] d \mathbf {w}}{\int \exp \left[ - \frac {1}{\sigma^ {2}} (\mathbf {X} \mathbf {w} - \mathbf {y}) ^ {T} (\mathbf {X} \mathbf {w} - \mathbf {y}) - \mathbf {w} ^ {T} \mathbf {w} \right] d \mathbf {w}} \\ = \frac {\int \mathbf {w} \exp \left[ - \frac {1}{\sigma^ {2}} \left(\mathbf {w} - (\mathbf {X} ^ {T} \mathbf {X} + \sigma^ {2} \mathbf {I} _ {D}) ^ {- 1} \mathbf {X} \mathbf {y}\right) ^ {T} (\mathbf {X} ^ {T} \mathbf {X} + \sigma^ {2} \mathbf {I} _ {D}) \left(\mathbf {w} - (\mathbf {X} ^ {T} \mathbf {X} + \sigma^ {2} \mathbf {I} _ {D}) ^ {- 1} \mathbf {X} \mathbf {y}\right) \right] d \mathbf {w}}{\int \exp \left[ - \frac {1}{\sigma^ {2}} \left(\mathbf {w} - (\mathbf {X} ^ {T} \mathbf {X} + \sigma^ {2} \mathbf {I} _ {D}) ^ {- 1} \mathbf {X} \mathbf {y}\right) ^ {T} (\mathbf {X} ^ {T} \mathbf {X} + {\sigma^ {2}} \mathbf {I} _ {D}) \left(\mathbf {w} - (\mathbf {X} ^ {T} \mathbf {X} + {\sigma^ {2}} \mathbf {I} _ {D}) ^ {- 1} \mathbf {X} \mathbf {y}\right) \right] d \mathbf {w}} \\ = \left(\mathbf {X} ^ {\intercal} \mathbf {X} + \sigma^ {2} \mathbf {I} _ {D}\right) ^ {- 1} \mathbf {X} ^ {\intercal} \mathbf {y} \tag {8} \\ \end{array}
$$
where  $\mathbf{X} = (\mathbf{x}_1^{\mathsf{T}},\dots,\mathbf{x}_{k - 1}^{\mathsf{T}})\in \mathbb{R}^{(k - 1)\times D}$  and  $\mathbf{y} = (y_{1},\ldots ,y_{k - 1})$
For most experiments, we study linear regression in  $D = 8$  dimensions with up to  $K = 16$  in-context examples and observation noise variance  $\sigma^2 = 0.25$ . Our base model is a transformer with the GPT2 architecture [15] with 8 layers, 128-dimensional embeddings, and 2 attention heads. We train with
the Adam optimizer [16] and a one-cycle triangle learning rate schedule [17] with  $50\%$  warmup. The base model is trained with batch size 256 for 500K training steps, though these hyperparameters are varied in our experiments. Specifically, for the experiments in Fig. 3 we sweep over 500K and 1M training steps to find the task diversity threshold as a function of training steps. For all other experiments, we sweep batch size 256 and 512 and for the experiments in Fig. 2, we do an additional batch size of 1024.
For the experiments in Fig. 6, we use a larger model with 12 layers, 256-dimensional embeddings, and 4 attention heads. This ensures that the model can learn to perform ridge regression in the  $D = 32$  setting. We also train the models on  $D = 8, 16, 24, 32$  with up to  $K = 2D$  in-context examples in each case. The observation noise variance is scaled at each dimension to keep the signal to noise ratio constant. So  $\sigma^2 = 0.5, 0.707, 0.866, 1$  for  $D = 8, 16, 24, 32$  respectively.
We always sweep over learning rates in  $\{0.0001, 0.0003, 0.001\}$  and choose the largest learning rate at which training is stable. For almost all experiments, the learning rate is 0.001, the main exception being the experiments in Fig. 6 where the learning rate is 0.0001.
For the 2M training step experiments in Fig. 4 and Fig. 10, the learning rate increases linearly up to 0.001 at  $250\mathrm{K}$  steps (i.e. using the same warmup as in the  $500\mathrm{K}$  step experiments) and then decreases linearly to 0 at 2M steps, consistent with our choice of decayed learning rate schedules. The experiments in Appendix F are performed at batch size 512 and following the  $500\mathrm{K}$  step one-cycle triangle learning rate schedule described above.
The implementation for this paper was done in JAX and all experiments were run on TPU v2-8 and v3-8 provided as part of the Google TPU Research Cloud program. Each training run takes approximately 4 v2-8 TPU hours.
C Support Figure 2
Fig. 8 provides an additional visualization for Fig. 2 upper middle but on a linear scale on the y-axis. In this figure it is clear that for all task diversities below the threshold—less that  $2^{14}$ —increasing the batch size aligns the PT with dMMSE, the optimal estimator for  $\mathcal{T}_{\mathrm{Pretrain}}$ . Thus below the task diversity threshold, the optimal estimator does indeed approach the Bayesian estimator with prior  $\mathcal{T}_{\mathrm{Pretrain}}$  when trained on more data and is thus unable to learn new tasks as discussed in Section 3.1.
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 of pretraining tasks. See Appendix C 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. 2 middle and right column, we probe this by increasing the batch size to increase the number of sequences per task seen during pretraining at each level of task diversity. To ensure that the phase transition and task diversity threshold we find is indeed a consequence of increasing the number of sequences per task and not merely just an effect of increasing batch size, we keep the batch size constant and increase the number of sequences per task in Fig. 3.
In Fig. 9 we provide additional evidence that the number of sequences per task and not batch size nor number of training steps is the primary factor that drives the behavior of the final model. Specifically, we show that the final behavior of the model is invariant to batch size or number of training steps if the number of training sequences per task is kept constant, at least within a reasonable range of batch size and number of training steps. To do so, at each level of pretraining task diversity, we compare our base model (batch size = 256, 500K training steps, red circles) to a model trained with batch size 512 (light blue squares) and batch size 128 (purple triangles). For either variant of the batch size, we adjust the number of training steps to keep the total number of pretraining sequences per task constant: 250K training steps for batch size 512 and 1M training steps for batch size 128. In all three cases, the behavior of the PT is identical as measured by  $\Delta_{\mathrm{PT,dMMSE}}^{\mathcal{T}}$  and  $\Delta_{\mathrm{PT,Ridge}}^{\mathcal{T}}$  for both  $\mathcal{T}_{\mathrm{Pretrain}}$  and  $\mathcal{T}_{\mathrm{True}}$  across all numbers of pretraining tasks.
Figure 9: 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 3. Each panel corresponds to the same panel in Fig. 3. See Appendix D for details.
E Small PT
We run our task diversity threshold finding experiment in a small PT with 64 dimensional embeddings, 4 layers, and 2 attention heads. The existence of a task diversity threshold reproduces in a Small PT: the PTs transition from becoming less like Ridge to becoming more like Ridge when trained on more data (see Fig. 10). We also note that the Small PT with a lower capacity has a lower threshold  $(\sim 2^{11.5}$  compared to  $\sim 2^{14.5}$  for the Base PT in Fig. 2) but also has lower overall performance (compare y-axis range to Fig. 2 lower right).
Figure 10: Existence of a task diversity threshold reproduces in small PT. Starting from batch size 256 and  $500\mathrm{K}$  training steps, we train on 2x data (i.e. examples per task) by doubling the batch size, and on 8x data by doubling batch size and quadrupling the number of training steps. We show distance to Ridge on new tasks (as in Fig. 2 lower right), and also plot distance to Ridge for models trained with infinite task diversity, where each training example samples a new task from the Gaussian distribution, as dashed lines (colors indicating batch size and training time). Points after the threshold  $(M \geq 2^{12})$  are shaded gray. Importantly, the threshold value doesn't change even if we go from 2x to 8x data suggesting that it isn't an artifact of underfitting.
Figure 11: The training loss barrier between a PT trained with  $M \gtrsim 2^{13}$  and a PT trained with  $M = \infty$  is small. At each  $M$ , including  $M = \infty$ , we train four small PTs, each with a different random seed for tasks. Each purple dot is an average over the 16 finite-infinite task model pairs at the given value of  $M$ , and each light blue dot is an average over the 6 infinite-infinite task model pairs. Error bars are standard deviations. See Appendix F 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 of tasks,  $M$ , as well as  $M = \infty$  (where each training example samples a new task from the Gaussian distribution), we train four small PTs, each with a different random seed for tasks. This means that all randomness (including model initialization, data, and observation noise) is shared across the four experiments, except for that governing the sampling of  $\mathbf{w}$ 's used for training. We then compare the finite  $M$  models to the  $M = \infty$  models, both in terms of the size of the training loss barrier along a path interpolating between the models, as well as the distance in weight space between the models.
We compute the training loss barrier between a finite  $M$  model and an infinite task model under the objective  $L_{\mathrm{Pretrain}}^T$  defined by the specific set of  $M$  tasks the finite task model was trained on. We follow the approach outlined in [31] where the barrier between two models is the loss at a model whose weights are the average of the two models, minus the average of the losses of the two models. Specifically, given two weight configurations  $w_{1}$  and  $w_{2}$  and a loss function  $L$ , the loss barrier is calculated as
$$
\text {l o s s} = L \left(\frac {w _ {1} + w _ {2}}{2}\right) - \frac {L \left(w _ {1}\right) + L \left(w _ {2}\right)}{2} \tag {9}
$$
The distance in weight space, in turn, is simply the  $L_{2}$  norm of the difference in model weights. For each  $M$  we average barriers and distances across all finite-infinite task pairs of models. We treat the analogous quantities computed between pairs of infinite task models as a baseline. See Fig. 11 and the Discussion section in the main text.
G Smoothed dMMSE
Figure 12: The PT outperforms the Smoothed dMMSE estimator with optimal smoothing. See Appendix G.
Here we consider the setting in which the discrete point prior over  $M$  tasks seen in pretraining is replaced with a mixture of  $M$  isotropic Gaussians with variance  $\epsilon^2$ . We call the estimator that does posterior predictive inference under this prior the Smoothed dMMSE estimator (sMMSE). In the limit of  $\epsilon \to 0$ , the sMMSE estimator becomes the dMMSE estimator. On the other hand, if  $\epsilon \to \infty$ , the prior will look like a Gaussian distribution but will have a variance much larger than the variance of  $\mathcal{T}_{\mathrm{True}}$ . In both edge cases, it won't perform well on  $\mathcal{T}_{\mathrm{True}}$ , and so there is an optimal  $\epsilon$  that achieves the best performance on  $\mathcal{T}_{\mathrm{True}}$ . In practice, we perform a search by simulation to determine the optimal  $\epsilon$ . We observe in Fig. 12 that for  $2^{13}$  tasks onward, the PT outperforms the optimally smoothed sMMSE on  $\mathcal{T}_{\mathrm{True}}$ .
To derive the sMMSE estimator, we want to calculate  $\hat{\mathbf{w}}_k = \mathbb{E}[\mathbf{w}\mid S_k]$  and apply Eq. (6).
$$
\begin{array}{l} p (\mathbf {w} | S _ {k}) = \frac {p (S _ {k} | \mathbf {w}) p (\mathbf {w})}{p (S _ {k})} \\ \propto \prod_ {j = 1} ^ {k - 1} p (y _ {j} | \mathbf {x} _ {j}, \mathbf {w}) \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2 \epsilon^ {2}} \left(\mathbf {w} - \mathbf {w} ^ {(i)}\right) ^ {2}\right) \\ \propto \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2 \sigma^ {2}} \sum_ {j = 1} ^ {k - 1} \left(\mathbf {w} ^ {\intercal} \mathbf {x} _ {j} - y _ {j}\right) ^ {2} - \frac {1}{2 \epsilon^ {2}} \left(\mathbf {w} - \mathbf {w} ^ {(i)}\right) ^ {2}\right) \\ \propto \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2} \mathbf {w} ^ {\intercal} \left(\frac {1}{\sigma^ {2}} \mathbf {X} ^ {\intercal} \mathbf {X} + \frac {1}{\epsilon^ {2}} \mathbf {I}\right) \mathbf {w} + \left(\frac {1}{\epsilon^ {2}} \mathbf {w} ^ {(i)} + \frac {1}{\sigma^ {2}} \mathbf {X} ^ {\intercal} \mathbf {y}\right) ^ {\intercal} \mathbf {w} - \frac {1}{2 \epsilon^ {2}} \left\| \mathbf {w} ^ {(i)} \right\| ^ {2}\right) \\ \end{array}
$$
We define,
$$
\mathbf {A} = \frac {1}{\sigma^ {2}} \mathbf {X} ^ {\intercal} \mathbf {X} + \frac {1}{\epsilon^ {2}} \mathbf {I}
$$
$$
\mathbf {J} _ {i} = \frac {1}{\epsilon^ {2}} \mathbf {w} ^ {(i)} + \frac {1}{\sigma^ {2}} \mathbf {X} ^ {\intercal} \mathbf {y}
$$
$$
\begin{array}{l} C = \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2 \epsilon^ {2}} \| \mathbf {w} ^ {(i)} \| ^ {2}\right) \int d \mathbf {w} \exp \left(- \frac {1}{2} \mathbf {w} ^ {\intercal} \mathbf {A} \mathbf {w} + \mathbf {J} _ {i} ^ {\intercal} \mathbf {w}\right) \\ = \sqrt {\frac {(2 \pi) ^ {D}}{| \mathbf {A} |}} \sum_ {i = 1} ^ {M} \mathbf {J} _ {i} ^ {\intercal} \mathbf {A} ^ {- 1} \mathbf {J} _ {i} \exp \left(- \frac {1}{2 \epsilon^ {2}} \| \mathbf {w} ^ {(i)} \| ^ {2}\right) \\ \end{array}
$$
so that,
$$
p (\mathbf {w} \mid S _ {k}) = \frac {1}{C} \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2} \mathbf {w} ^ {\intercal} \mathbf {A} \mathbf {w} + \mathbf {J} _ {i} ^ {\intercal} \mathbf {w} - \frac {1}{2 \epsilon^ {2}} \| \mathbf {w} ^ {(i)} \| ^ {2}\right)
$$
Plugging back to Eq. (6),
$$
\begin{array}{l} \hat {\mathbf {w}} _ {k} ^ {\mathrm {s M M S E}} = \frac {1}{C} \sum_ {i = 1} ^ {M} \exp \left(- \frac {1}{2 \epsilon^ {2}} \| \mathbf {w} ^ {(i)} \| ^ {2}\right) \int d \mathbf {w} \mathbf {w} \exp \left(- \frac {1}{2} \mathbf {w} ^ {\intercal} \mathbf {A} \mathbf {w} + \mathbf {J} _ {i} ^ {\intercal} \mathbf {w}\right) \\ = \frac {1}{C} \sum_ {i = 1} ^ {M} \sqrt {\frac {(2 \pi) ^ {D}}{| \mathbf {A} |}} \mathbf {A} ^ {- 1} \mathbf {J} _ {i} \exp \left(\frac {1}{2} \mathbf {J} _ {i} ^ {\intercal} \mathbf {A} ^ {- 1} \mathbf {J} _ {i} - \frac {1}{2 \epsilon^ {2}} \| \mathbf {w} ^ {(i)} \| ^ {2}\right) \\ \equiv \sum_ {i = 1} ^ {M} \beta_ {i} \mathbf {A} ^ {- 1} \mathbf {J} _ {i} \\ \end{array}
$$
where  $\beta = \mathrm{softmax}(\tilde{\beta})$  and  $\tilde{\beta}_i\equiv \frac{1}{2}\mathbf{J}_i^\top \mathbf{A}^{-1}\mathbf{J}_i - \frac{\|\mathbf{w}^{(i)}\|^2}{2\epsilon^2}$
H Supporting figures
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 is the crossover point between  $\Delta_{\mathrm{PT,Ridge}}^{\overline{T}_{\mathrm{True}}}$  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 to those of Ridge, as quantified by  $\Delta_{\mathrm{PT,Ridge}}^{\overline{T}_{\mathrm{True}}}$  at 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 no discernible effect on the threshold (bottom) for a base PT.
Figure 14: Task diversity threshold for tasks sampled from a Laplace prior. We run our task diversity threshold finding experiment in a setting where each element of each task regression vector is 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 on new tasks by training for 500K steps with batch size 512 on a 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 a task diversity threshold between  $2^{14}$  and  $2^{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.