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

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 $\textit{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 $\textit{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 $\textit{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 $\textit{all tasks}$, including those not seen during pretraining. Thus, when pretrained on data with task diversity greater than the threshold, transformers $\textit{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. Code is available at https://github.com/mansheej/icl-task-diversity.

Paper

References (33)

Scroll for more · 21 remaining

Similar papers

Peer review

Reviewer CEjq8/10 · confidence 4/52023-06-11

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.

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. 2. 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. 3. The experiments are rigorous. 4. 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. 2. 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.

Questions

N/A

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Yes

Reviewer Uwh77/10 · confidence 4/52023-06-18

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.

Strengths

1. The paper is well-written and introduces all the necessary background needed for understanding. 2. 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. 3. The empirical study is well controlled and supports the conclusions made.

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. 2. 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. [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. [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. [3] K. Hsu, S. Levine, and C. Finn, “Unsupervised Learning via Meta-Learning.” arXiv, Mar. 21, 2019. doi: 10.48550/arXiv.1810.02334. [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.

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. - 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." 2. 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. 3. 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)? - 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. - Similarly, the paper shows that one (meta-)learning algorithm (next token prediction with Transformers) is superior to the Bayesian dMMSE algorithm. - 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)$. - Are these analogies (fairly) the same, or is there a fundamental difference?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

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. - 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. 2. 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] [1] M. Garnelo et al., “Conditional Neural Processes.” arXiv, Jul. 04, 2018. doi: 10.48550/arXiv.1807.01613.

Reviewer NkJG7/10 · confidence 4/52023-06-30

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. Surprisingly, 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.

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!

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 teacher weights, if the Transformer is large enough simply learn and remember all needed weights ("features") and implement dMMSE. Now, when the capacity is not sufficient anymore it falls back to ridge regression which happens to generalize better. In 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. Although I find the paper still interesting, for me to give a higher score various questions and experiments are needed. I have a couple of open ended questions on which I would like to hear your thoughts. See questions. Minor things: One 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.

Questions

Unfortunately, to give a higher score I feel that the paper should go in a slightly different direction. Questions for a higher score: Which 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? Open ended: 1) 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. The goal would be to understand if one can push towards the generalizing solution with the same amount of data. 2) 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 at play here. 3) 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

Limitations are obvious and mentioned in the discussion.

Reviewer nyxK6/10 · confidence 4/52023-07-05

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).

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). 2. Find that weight norm improves this diversity threshold in the particular setup. 3. A clear and concise writing.

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? 2. 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. 3. 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. 4. Would be more complete if analysis of hyper parameters can be extended to the number of layers too.

Questions

1. Are these transitions hold for different distributions than Guassian (Laplace, etc)? 2. 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? 3. 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. **Paper Summary** I 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors could discuss the possibility of this transition might not be hapenning in the case of enough iteration and samples per task.

Reviewer vcDA4/10 · confidence 5/52023-07-07

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.

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. 2) 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. 3) 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. 1) 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. 2) 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)?

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. 2) For the models that converged to sMMSE, how many steps does it take? Is it the same number of steps across task counts? 3) (Minor) In Figure 2, is the upper right figure the same as the lower right figure?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors do not have a separate limitations section but discuss alternative explanations for experiments throughout the paper.

Reviewer NkJG2023-08-10

Thank you!

I have read the rebuttal, the other opinions and will increase my score, voting for acceptance. Thank you!

Reviewer vcDA2023-08-12

Thank you, I appreciate the new experiments! I find the order of learning Ridge than dMMSE important and I hope it gets adequately situated within the discussion of the threshold. I am currently reconsidering my score, and I wanted to ask some questions before I do so. The results suggest that the task diversity threshold for Small PT is $2^{11.5}$ and the task diversity threshold for PT is approximately $2^{14.5}$. From this, I'd imagine that as you increase model capacity, the task diversity threshold increases. I believe for large $M$, the models are not in an over-parameterized regime [1] since the model can not solve the training distribution (evidenced by not achieving dMMSE performance). **Candidate Explanation**: The "task diversity threshold" is the model only having the capacity to represent Ridge and dMMSE for task counts below the threshold and the model not being able to express dMMSE for task counts above the threshold. As for the improved out-of-distribution generalization, this is complicated by the model learning Ridge first and then dMMSE for medium-sized $M$. As such, the model learning Ridge does not reflect out-of-distribution generalization as much as it reflects the model not generalizing in-distribution (due to capacity constraints). I was curious whether the experiments shed any insight into whether this explanation is feasible. It would help me understand how this work expands our current understanding of generalization in ML, especially with respect to the standard bias-variance tradeoff (or the under-parameterized regime of double descent). [1] Reconciling modern machine learning practice and the bias-variance trade-off

Authorsrebuttal2023-08-15

We are glad you appreciate the new experiments, and we thank you for your important questions, which we attempt to answer. The reviewer correctly notes that as model capacity (at least as measured in number of parameters) increases, so does the task diversity threshold $M^*$, especially when the model capacity is small to begin with. Now when $M > M^*$ for any given model capacity, our simulations indicate that our model does not find the dMMSE solution which would minimize in-distribution error $L^{T_\text{Pretrain}}$. There could be two main reasons for this: (1) there is a model capacity/expressivity issue, namely there is *no* configuration of model weights that can achieve the dMMSE estimator when $M>M^*$, or (2) there is a trainability issue, namely, there *is* a configuration of model weights that could achieve dMMSE, but its basin of attraction under SGD is so small that SGD starting from a random initial condition does not find it, and instead SGD typically is driven to the training error minimum of a basin containing a *different* solution. The reviewer suggests (1) as a possible explanation. Our current experiments may be consistent with explanation (1), but also do not rule out explanation (2) as a possibility. To rule out explanation (2) one would have to find the dMMSE solution without SGD, and it is not clear how to do that. This is characteristic of all failures of machine learning: when it fails we do not know if the failure is due to expressivity alone or due to trainability via SGD. That being said, even in a setting in which the model cannot fit dMMSE (regardless of whether the failure is due to expressivity or trainability), there exists an infinite set of underfitting (in-distribution) solutions which achieve lower in-distribution training error. Apriori, it is unlikely that, at a given threshold value of $M$, all of these underfitting solutions become inaccessible. However, our experiments show that: when the tasks are drawn from a Gaussian, beyond the task diversity threshold, amongst the infinite set of all possible underfitting solutions, _there is a preference for ridge regression as **the** underfitting solution_, and the PT reaches $L^{T_\text{True}}$ performance on par with models trained in the infinite task setting. When the tasks are drawn from a Laplace distribution, in turn, beyond the task diversity threshold, there is a preference for the Bayes-optimal MMSE estimator under a Laplace prior. **Hence, the model doesn’t just learn any underfitting solution. It learns the solution that optimizes error for the underlying generative model of pretraining tasks. The fact that it performs on par with the infinite task setting solution, despite having access only to a finite set of tasks, means that it is, remarkably, correctly generalizing out of the training distribution, even though it is not generalizing in-distribution.** This an interesting example of how the failure to generalize to a limited diversity in-distribution setting does not mean one cannot generalize to a more important, diverse out-of-distribution setting. This setting is not often considered in classical learning theory. Finally, we note that the reviewer’s description of the model’s “order of learning Ridge than dMMSE” is indeed important; in our revision we will be giving our rebuttal Fig. 2 strong prominence in the main paper. However, we would like to make this statement more precise. Fig. 2 does show that when $M$ is below the task diversity threshold, as measured by $\Delta_\text{PT, Ridge}$ the model does get closer to ridge over the course of training, and then moves farther away towards dMMSE. However, at these intermediate values of $M$ the lowest value of $\Delta_\text{PT, Ridge}$ achieved by the PT over the course of training is still considerably higher than that achieved by the PT in the infinite task setting. This means the PT _does not_ completely learn Ridge before starting to overfit at these intermediate values of $M$.

Reviewer vcDA2023-08-15

Thank you for this further clarification. In particular, I agree with the learning-theoretic observation as provided in the rebuttal. I did not get this message from the original paper; I believe the paper would benefit from this nuanced discussion of how SGD + Transformers are incapable of fitting the optimal in-distribution performance and it is interesting that of all possible under-fitting solutions, the model converges to the one that generalizes out-of-distribution. This new message of the work should also be situated within the context of Simplicity Bias literature, which discusses how among multiple possible solutions, models tend to converge to the simplest one, and it seems that Ridge is a simpler solution in some sense since it's discovered earlier in training and able to be learnt by smaller Transformers. Due to this extended discussion/clarifications about the results of the paper and the insightful clarifications provided by the reviewers, I am raising my score from a 3 to a 4. However, I maintain that showing the model converges to the Bayes-optimal solution for high task counts is a narrow contribution specific to this setup and is reliant on the 1) simplicity bias of getting closer to ridge first and 2) being in the under-parameterized regime.

Reviewer Uwh72023-08-12

Thank you for the clarifications and additional results. One thing that still concerns me is the argumentation about the shown behavior to be an "emergent" one. It still seems a bit misleading to me. As also noted by vcDA, this might be simply the inability of the model to fit the optimal dMMSE solution for large $M$s due to the limited capacity. This is also in line with the new results provided in Fig.1(B), where larger models can fit dMMSE solution for larger $M$s => an increased task diversity threshold. At the same time, I believe that the experimental setting considered in this work and the results provided through rigorous experimentation would be a valuable contribution to the field, regardless of their current and future interpretation. For example, even if it is a "simple underfitting," the fact that the model converges to the Ridge solution and not any other solution could be an interesting insight. I suggest that authors use "emergence" more carefully and discuss other alternative explanations clearly and explicitly. I am reconsidering my final grade, and I will follow the further discussion for that.

Authorsrebuttal2023-08-15

Thank you for the follow-up comments! We’ll be very clear in our paper about how we define emergence, which is that, when the number of tasks is above the task diversity threshold, although the PT model is trained on samples from $T_{\text{Pretrain}}$: 1. The PT achieves better performance on samples from $T_{\text{True}}$ than the Bayes-optimal estimator with prior $T_{\text{Pretrain}}$. 2. Moreover, above and beyond (1), the PT achieves strong (near-optimal) performance on samples from $T_{\text{True}}$, despite the fact that $T_{\text{True}}$ and $T_{\text{Pretrain}}$ are different. We certainly agree that it’s important to use this term carefully, that is adhering to a definition while also explaining possible causes, including the possibility that either the model is capacity-constrained and not expressive enough to implement the dMMSE solution, or the model is expressive enough to express dMMSE but this solution has such a small basin of attraction that SGD cannot find it. However, regardless of why SGD for the PT cannot find the dMMSE solution when $M$ is bigger than the task diversity threshold, **we believe that a fundamental finding of our paper is that, out of the infinitely many in-distribution underperforming solutions for $T_{\text{Pretrain}}$ that the PT could find, it remarkably finds the Ridge solution, which is optimal for the out-of-training distribution $T_{\text{True}}$**. It is precisely this behavior that we are calling emergent (as defined in joint points 1 and 2 above). This behavior is surprising and emergent precisely because nothing in $T_{\text{Pretrain}}$ necessarily pre-ordained this. We will be very careful to elucidate that emergence in this context means this and nothing more. Finally, a way to think about the dependence of the task diversity threshold on model capacity is as follows: think about a 2D phase diagram where one axis is task diversity (i.e. $M$) and the other axis is model capacity ($C$). At each point in the 2D phase diagram, we can ask a question: if we train a model with capacity $C$ and task diversity $M$ on more data (either more iterations as in Rebuttal Figs. 1A and 2A right and paper Fig. 3, or larger batch size as in Rebuttal Fig. 1A and paper Fig. 2) then, does the PT become more like Ridge or less like Ridge (and closer to dMMSE)? The answer to this binary question yields a phase transition boundary in the 2D phase diagram where on one side of the boundary (i.e. small $M$ and large $C$) the PT becomes more like dMMSE, and on the other side of the phase boundary, (i.e. large $M$ and small $C$) the PT becomes more like Ridge. We have been implicitly delineating this phase boundary by finely sampling $M$ and coarsely sampling $C$. It is entirely expected that the phase boundary should depend jointly on $M$ and $C$, hence the task diversity threshold in $M$ should grow modestly at least with $C$. We do not view this as a negative result of our work, but rather the elucidation of an important phase transition in a two parameter joint space of task diversity and model capacity.

Reviewer Uwh72023-08-17

I thank the authors for the provided discussion, which clarifies my questions. I generally agree with the concerns raised by vcDA. However, I believe the contribution can still be valuable enough for the community and increase my grade to 7, advocating for acceptance (provided the authors include more explicit discussion about the emergence and under-parametrization in the camera-ready version).

Reviewer nyxK2023-08-15

thank you!

Thank you for the detailed response and extra experiments. The transition seemed to me really interesting and further experiments that clarify this holds for many data points addressed my concerns. I increased my point as a result.

Reviewer CEjq2023-08-18

Thank you to the authors for your detailed response, and for your suggestion of the insightful reference. I also appreciate the authors' and other reviewers' thorough discussions. Having read these, I have decided to maintain my positive score. Nevertheless I would encourage the authors to add the new experiments and discussion, especially as it relates to whether the task diversity affect is an artifact of undertraining, to the updated paper.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC