A Theoretical Analysis of In-context Task Retrieval and Learning

We would like to express our gratitude to the reviewer for dedicating your valuable time to provide constructive feedback on our work. We hope our responses have addressed your concerns and answered your questions. If our responses have resolved your concerns, we kindly request you to consider increasing your score.

---

(1) “I would expect new studies” & “It’s not clear to me what the main takeaway of this paper”

Our paper makes two major novel contributions.

Firstly, we are the pioneering work that formally defines and analyzes task learning and retrieval. This can assist future researchers in reconciling seemingly conflicting observations made with in-context learning in the past, as they can be interpreted as non-conflicting results within the two different regimes.

Secondly, we are the first to precisely predict two non-monotonic performance patterns of ICL in real-world LLMs as the number of samples increases: the U-shaped curve we predicted in the original submission, and the flipped one we presented in the rebuttal. The U-shaped curve was confirmed by our additional experiments. The flipped curve was initially presented in the original GPT3 paper (and later was re-confirmed by Xie et al.’s work too!) and explained via experiment under our framework. Experiments are described in “supplementary material”.

These two contributions lead us to strongly believe that our paper offers substantial contributions to the research community, particularly theoretical analysis of ICL, and to some extent, practical LLM research.

---

(2) (paraphrased) “data generation process should be more realistic” & “It’s not clear how the generation process is related to the real-world data”

We deliberately selected the linear regression mixture model, which, as the reviewer correctly pointed out, is slightly less realistic (non-discrete) than the HMM model. Our choice was motivated by the model's analytical tractability, which provides a precise theoretical understanding of how task retrieval and task learning operate.

Numerous theoretical studies adopt the same linear regression mixture model inluding recent works [1-5] (see these papers in “To AC and All Reviewers” due to the character's length). Although the linear regression mixture model deviates from the actual pretraining data model, the insights it offers may still predict real-world phenomena.

---

(3) “The generation process of the pretraining data is the same as the downstream task”

We do not make such an assumption. Although we assume that the pretraining task parameters are drawn from a probability distribution whose support spans the entire task parameter space, this does not imply that every task is included in a finite pretraining dataset. In fact, since our task parameters are continuous random variables, with probability 1, the downstream task is not included in a finite pretraining dataset.

---

(4) “How is the definition of the two kinds of tasks relevant in the real-world scenario?”

Our work is largely motivated by various observations made with real-world LLMs that appear to be conflicting. Here are the most prominent examples:

Paper "Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?" demonstrates that in-context learning with random labels sometimes outperforms in-context learning with true labels. How can this be explained?

The GPT3 paper and the paper by Xie et al. both show that zero-shot performance is often higher than few-shot performance in various tasks. Notably, these findings are all based on ground truth labels. How can this be explained?

From a more traditional supervised-learning perspective, these two observations seem counterintuitive. The general expectation is that the more samples provided, the better the performance should be. Indeed, most theoretical ICL papers assume this setting and prove that ICL performance should monotonically increase as the number of samples increases. We refer to this setting as "task learning".

To explain these phenomena, we introduce "task retrieval", which happens when in-context examples are drawn from a different task than the target task (This can also be viewed as a form of concept shift, using distribution shift terminology). This might happen with label noise or when there is a subtle mismatch between the way labels are created for in-context samples versus evaluation samples. In this regime, we can rigorously show that when the number of samples is relatively small, the performance may not be monotonic.

---

(5) “Empirically, we do not observe that having more examples hurt the performance.”

In response to your insightful suggestion, we conducted experiments with GPT-4 and discovered that the U-shaped loss curve is observable in practice. The experiment arise described in “supplemental material”. Another related non-monotonic performance was reported in the original GPT-3 paper and reconfirmed by the paper by Xie et al. (please refer to Figure 7 therein).

Thank you for your strong support of our work!

We also would like to express our sincere gratitude to the reviewer for dedicating your valuable time to provide constructive feedback on our work. We added experiments to show the U-shaped loss curve actually exists in real-world LLM (GPT-4) and our framework can also be used to explain a flipped U-shaped loss curve researchers have observed in real-world LLM. We share the experiments in "supplementary material". We kindly request you to consider increasing your score and championing our paper if these new experiments significantly improve the quality and confidence of this paper.

---

(1) “High level proof idea for lemma, theory”

We will incorporate additional sentences to describe the high-level idea behind the proof.

---

(2) “Texts and figures are overlapped in Page 8”

Thanks! fixed.

We would like to express our gratitude to the reviewer for dedicating your valuable time to provide constructive feedback on our work. We hope our responses have addressed your concerns and answered your questions. If our responses have resolved your concerns, we kindly request you to consider increasing your score.

---

(1) “Explain the importance of studying this model”

We selected the linear regression mixture model, which is slightly less realistic (non-discrete) than the HMM model. Our choice was motivated by the model's analytical tractability, which provides a precise theoretical understanding of how task retrieval and task learning operate.

Numerous theoretical studies adopt the same linear regression mixture model. Notable recent works include [1] through [5]. All theoretical work necessitates some form of simplified model to comprehend complex real-world phenomena. Although the linear regression mixture model deviates from the actual pretraining data model, the insights it offers may still predict real-world phenomena.

[1] What can transformers learn in-context? a case study of simple function classes

[2] What learning ¨ algorithm is in-context learning? investigations with linear models

[3] Transformers as algorithms: Generalization and implicit model selection in in-context learning

[4] Transformers learn in-context by gradient descent

[5] Transformers can optimally learn regression mixture models

---

(2) “Some area of works that are related are left”

Thank you for sharing these papers! We will incorporate these papers into the related works section.

---

(3) “The dependence of U shaped risk bound on component reweighting and shifting is not discussed properly.”

We apologize if our current explanation is unclear. We will revise it to improve clarity.

The U-shape is observed under conditions of low prior task variance, and we aim to retrieve a prior task center that is closest to the in-context task. When k is small, due to the low prior task variance, component reweighting takes effect faster than component shifting. Thus, when all prior centers are nearly unchanged, the task center closest to the in-context task receives a high weight close to 1, thereby reducing task retrieval loss. However, as k continues to increase and becomes large, component shifting causes all posterior task centers to shift towards the in-context task. Therefore, when the in-context task differs from the target task, this results in an increase in task retrieval loss.

---

(4) “Connect the U shaped with the bias-variance tradeoff.”

Our U-shaped curve can not be directly explained by the bias-variance tradeoff. This is because the bias-variance tradeoff is for increasing model complexity, while our curve is for increasing the number of samples. Under the standard supervised learning setting, our U-shaped curve is not observed. Instead, our task retrieval setting is closely related to the distribution shift setting, particularly the concept shift.

---

(5) “The novelty in deriving the posterior distributions for the mixture of Gaussian distribution is unclear to me”

We have not claimed novelty in the calculation of the posterior distribution (although we believe it forms part of our contribution, given the extensive calculations required). Instead, our paper presents two major contributions in terms of novelty.

Firstly, we are the pioneering work that formally defines and analyzes the two regimes of in-context learning: task learning and retrieval. This can assist future researchers in reconciling seemingly conflicting observations made with in-context learning in the past, as they can be interpreted as non-conflicting results within the two distinct regimes.

Secondly, we are the first to precisely predict two non-monotonic performance patterns of ICL in real-world LLMs as the number of samples increases: the U-shaped curve we presented in the original submission, and the flipped one we presented in the rebuttal. The U-shaped curve was first predicted by our theory and then confirmed by our additional experiment with GPT4. The flipped curve was initially presented in the original GPT3 paper (and later was re-confirmed by Xie et al.’s work too!) and explained via experiment under our framework. The experiments are provided in “supplementary material”

These two contributions lead us to strongly believe that our paper offers substantial contributions to the research community, particularly those focused on the theoretical analysis of ICL, and to some extent, practical LLM research.

Thank you for your strong support of our work!

We also would like to express our sincere gratitude to the reviewer for dedicating your valuable time to provide constructive feedback on our work. We hope our responses have addressed your concerns and answered your questions. If our responses have resolved your concerns, we kindly request you to consider increasing your score and championing our paper.

---

(1) “the similarity to real-world settings is not quite accurate, as the setting studied in the paper is quite different from in-context learning in NLP tasks” & “assumptions in Section 3.3 are restrictive, and the scope of the theoretical analysis is limited to a very specific setup”

We deliberately selected the linear regression mixture model, which, as the reviewer correctly pointed out, is slightly less realistic (non-discrete) than the HMM model. Our choice was motivated by the model's analytical tractability, which provides a precise theoretical understanding of how task retrieval and task learning operate. Numerous theoretical studies adopt the same linear regression mixture model. Notable recent works include [1] through [5].

All theoretical work necessitates some form of simplified model and strong assumptions to comprehend complex real-world phenomena. Although the linear regression mixture model and our strong assumptions may not look very realistic, the analysis we performed with this model still predict real-world phenomena.

[1] What can transformers learn in-context? a case study of simple function classes [2] What learning ¨ algorithm is in-context learning? investigations with linear models [3] Transformers as algorithms: Generalization and implicit model selection in in-context learning [4] Transformers learn in-context by gradient descent [5] Transformers can optimally learn regression mixture models

---

(2) “A highly expressive F can be viewed as K separate models F0, . . . , FK−1, where Fk takes exactly 2k + 1 tokens as input. Can the authors provide more explanation regarding this statement?"

F is a function that takes variable-length input. Thus, F can be fully characterized by specifying F's behavior for each input length. That is,

F = F_0 * 1(# of in-context examples = 0) + F_1 * 1(# of in-context examples = 1) + …

Now, consider a sequence of tokens [x1, y1, x2, y2]. Pretraining can be viewed as

min_F L(F([x1]),y1) + L(F([x1,y1,x2]),y2).

By the above input-length decomposition, this is equivalent to

min_{F_0, F_1} L(F_0([x1]),y1) + L(F_1([x1,y1,x2]),y2).

Thus, this becomes two separate optimization problems. Now, if F is highly expressive, then F_0 and F_1 are also highly expressive. Thus, each F_i can become the Bayes-optimal solution.

---

(3) “The code seems to suggest the architecture used is GPT-2, and NanoGPT, but these details should be provided in the paper.”

The code is constructed on the GPT2Model from the transformers package supported by Hugging Face. The number of attention heads is set to 8, the depth is set to 10, and the embedding dimension is set to 1024. We use the Adam optimizer with a learning rate and weight decay both set to 0.00001. The batch size is set to 256. We train the model for 3 epochs, with each epoch consisting of 10,000 steps.

---

(4) “Figures 6,7, and 8 are central results for the paper, and are not quite accessible for readers. It would be beneficial to add more experiments in the main paper,” & “text overlaps in Figure 5”

Thank you for your insightful suggestions! We will incorporate these changes into our revision.

Thank you for your strong support of our work!

We also would like to express our sincere gratitude to the reviewer for dedicating your valuable time to provide constructive feedback on our work. We hope our responses have addressed your concerns and answered your questions. If there are any remaining questions, please do not hesitate to let us know. If our responses have resolved your concerns, we kindly request you to consider increasing your score and championing our paper.

---

(1) “The generative modeling makes strong simplifying assumptions like Gaussian distributions that limits applicability to complex real-world textual data. Expanding the analysis to more realistic data distributions will strengthen it. Or perhaps even showing that any of their conclusions hold in LLMs, for example does the u shaped phenomena occur in LLMs?”

In response to your valuable suggestion, we conducted additional experiments to demonstrate that the U-shaped curve can indeed be observed with a real-world LLM (GPT-4). Please refer to Section 1 of the supplemental file for details.

Moreover, we demonstrate that under a slightly different setting, our task learning analysis can explain another intriguing phenomenon, which we term the 'flipped U curve' – a pattern where the loss increases and then decreases as the number of samples increases. This effectively explains the phenomenon reported in the original GPT-3 paper as well as in subsequent research. Please refer to Section 2 of the supplemental file for details.

Thus, even if our linear regression mixture model deviates from the actual pretraining data model, the insights it offers can predict real-world phenomena!

---

(2) ”optimal retrieval remains unspecified”

The question essentially requires finding the minimizer of $f(k) = a \exp(-bk) + c \exp(-dk^{1/2}) + e k^2$, which is equivalent to solving $f’(k)=0$. The author believes that there is no closed-form solution for this.

---

(3) “Evaluating other loss functions (Cross entropy)”

Thank you for the excellent suggestion. We are currently exploring it; however, it's challenging to find a suitable prior for classification that is akin to the Gaussian mixture used in the linear regression setting. The Gaussian mixture provides flexibility for mathematical calculation and exploration of the prior's properties. We have added evaluations on GPT-4 which is trained with cross-entropy to reproduce the U-shaped loss curve, in the Section 1 of the supplemental file.