Paper: A Theoretical Analysis of In-context Task Retrieval and Learning (ICLR-2024)

Reviewer comment (Theory):
Firstly, it’s a gaussian model, so it’s very different from the discrete case of NLP. Because it’s not discrete, it is even less realistic than the HMM model used by Xie et al.

Author response:
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.

Reviewer comment (Novelty):
It’s not clear to me how the definition of the two kinds of tasks is relevant in the real-world scenario.

Author response:
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.

Reviewer comment (Evaluation):
Empirically, we do not observe that having more examples hurt the performance.

Author response:
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).

Reviewer comment (Experiments):
It’s not clear how the generation process is related to the real-world data.

Author response:
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.

Reviewer comment (Theory):
I observed that Lemma 1 has been used a lot in the manuscript and is lack of proof. Please at least provide high-level idea of the proofing before using it.

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

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

Reviewer comment (Presentation):
Texts and figures are overlapped in Page 8.

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

Thanks! fixed.

Reviewer comment (Theory):
Similarly, for Theorem 3,4 and Lemma 5, please provide at least few sentences about proofing instead of pointing the appendix.

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

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

Reviewer comment (Novelty):
The generative model is simplistic as it is limited to linear Gaussian mixture. In contrast HMM based ICL is already studied in other works already mentioned in the paper. The authors should explain the importance of studying this model.

Author response:
(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

Reviewer comment (Novelty):
The authors leave out some area of works that are related, e.g. meta-learning and retrieval augmented learning. A few recent examples of the latter are -  'A Statistical Perspective on Retrieval-Based Models' by Basu et al., 'Generalization and stability in in-context learning' Li  et al.

Author response:
(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.

Reviewer comment (Theory):
The dependence of U shaped risk bound on component reweighting and shifting is not discussed properly. Furtheremore, U shape is also observed through the bias-variance tradeoff in optimization literature. The authors don't connect the U shaped mention here with the bias-variance tradeoff.

Author response:
(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.

Reviewer comment (Novelty):
The novelty in deriving the posterior distributions for the mixture of Gaussian distribution is unclear to me.

Author response:
(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.

Reviewer comment (Experiments):
The claim regarding 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 with practical LLMs.

Author response:
(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

Reviewer comment (Theory):
The assumptions in Section 3.3 are restrictive, and the scope of the theoretical analysis is limited to a very specific setup based on a particular generative model.

Author response:
(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

Reviewer comment (Theory):
"A highly expressive F can be viewed as K separate models F0, . . . , FK−1, where Fk takes exactly 2k + 1 tokens as input. Thus, pretraining can be decomposed into K separate optimization problems." Can the authors provide more explanation regarding this statement?

Author response:
(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.

Reviewer comment (Reproducibility):
Can the authors provide more details about the Transformer? The code seems to suggest the architecture used is GPT-2, and NanoGPT, but these details should be provided in the paper.

Author response:
(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.

Reviewer comment (Presentation):
The experimental results provided in Appendix A with 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, and move some of the analysis to the supplementary material.

Author response:
(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.

Reviewer comment (Presentation):
There are typographical issues and formatting issues (e.g. text overlaps in Figure 5) that should be fixed.

Author response:
(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.

Reviewer comment (Experiments):
The criteria for determining the number of examples k for optimal retrieval remains unspecified. Providing more precise theoretical or empirical guidance could improve utility.

Author response:
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.

Reviewer comment (Evaluation):
Risk is evaluated using squared loss. Evaluating other loss functions (Cross entropy) could expand usefulness across different domains.

Author response:
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.