The Learnability of In-Context Learning

Summary

The authors propose a PAC framework to analyze the expressiveness power of in-context learning in a finite sample complexity scheme. The framework consists of largely two parts, first is the initial pretraining of the next token prediction phase, and the second is the in-context learning phase. Regarding the pertaining data distribution as a mixture of latent tasks. authors provide an understanding that polynomial sample complexity guarantees in-context learnability.

Strengths

- Define a PAC framework to investigate in-context learning - Reduce the sample complexity needed for in-context learning to polynomial complexity, which is confined to infinite samples in previous work

Weaknesses

- The number of the model parameter is also a crucial factor for in-context learning, but it seems that there is not enough handling for model complexity.

Questions

- In line 143, the authors say they aim to provide an analysis of in-context learnability of “large model”. With the specified sample complexity, what is the guideline for that “large model”?

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

- Lack of experimental evidence to support the theory (at least experiments on synthetic data would be helpful)

Summary

This paper attempts to formalise the few-shot in-context learning phenomenon observed in large language models. To that end, they make a set of assumptions about the underlying data-generating distributions, pretrained models, etc and try to formalise the task of in-context learning in the PAC framework. The main empirical phenomenon that is the subject of study is the ability of LLMs to predict a query input accurately after receiving a sequence of pairs of inputs and labels. Broadly, they prove that if the distributions over strings for pretraining and providing in-context examples are defined in a certain way and satisfy a set of assumptions, then for large enough k, f(y | x_1, y_1, …, x_k, y_k) will be approximately correct with high probability (in the PAC sense) where f is the language model. In their framework, they assume that the target task during in-context learning is part of the pretraining distribution based on prior observations in empirical works. The hypothesis class for pretraining is a set of mixture distributions of downstream distributions. The downstream distributions can be seen as a distribution over pairs of inputs and labels. The sampling from the pretraining distribution can be decomposed into sampling a task from a prior distribution over the set of tasks and then sampling from the respective downstream distribution corresponding to the task. They assume that we have an accurate language model (in the PAC sense) where the error with the target distribution is bounded in terms of TVD of the conditional next word distribution -- based on the fact that we have LMs that are accurate in modelling the underlying distribution. Given such an accurate probabilistic LM and inputs from such distribution, they show that given a prompt, for large enough k, the error for f(y | x_1, y_1, …, x_k, y_k) with respect to the Bayes optimal predictor is less than ε with high probability.

Strengths

(S1) I think the formalisation is interesting to some extent, and their analysis provides some intuition as to why few-shot learning, as observed in GPT-3, works to some degree. I think the key idea is to use the fact that the ratio P_A(p) and P_B(p) converge to zero as k gets larger, where B is the target task, and A is a different task. This helps distinguish the tasks and improves the margin between the correct label and incorrect label for the LM as the size of the prompt goes larger. Another interesting part about their framework is that they treat it as a distribution modelling task, unlike some other recent efforts to formalise in-context learning [2, 3] and derive sample complexities [1]. They have taken a more physics-like approach, where they have assumed and defined the models and data-generating distributions in a certain way based on empirical observations and tried to work out why prompts could lead to correct answers if the pretraining data was not of that form. The four assumptions are not completely unrealistic. At least within their framework, it is somewhat clear why flipping some labels is not as detrimental as one would expect. (S2) In-context learning has been an intriguing phenomenon, with multiple works seeking to understand it. I think this paper takes a step towards that and could be useful to other researchers working in this area. (S3) The paper is well written, easy to follow, and the arguments are clearly presented.

Weaknesses

While the analysis and the framework are interesting, I think there are a few issues with the framework that limits its applicability in helping us understand in-context learning. (W1) The way the pretraining distribution is defined is not necessarily reflective of the real-world data. To be clear, it need not always be and depends on the nature of the theoretical work. For instance, the setup [1-3] is quite simplified, but they still allow us to train Transformers on those specific tasks and test how well the theory predicts the sample complexities in the simplified setting. In the case of this paper, the theory seems to be an attempt to directly model in-context learning in the real-world scenario, and the way the pretraining distribution is defined seems a bit detached from real-world text. Additionally, if my understanding is correct (correct me if I am wrong), each example from the pretraining distribution seems to be an input example followed by a label from a downstream task. This seems a bit far from the way LLMs are trained as well. To some extent, it seems like the pretraining and downstream distributions are designed to satisfy some properties favourable for the end result but not necessarily reflective of the real data. I understand that it is not possible to precisely define real-world data, but given the goal of the paper, it seems like the value of results does depend on how well the data-generating distribution reflects real-world data. (W2) It seems like the sample complexities are not necessarily predictive of the number of few-shot examples needed for in-context learning. I think it stems from the way the framework is defined and the disconnect with practice. Standard learning theoretic frameworks such as PAC models are entirely formal abstractions and do not contain assumptions about the real world (apart from train and test being the same distribution) -- hence allowing worst-case sample complexities to be meaningful. It also allows us to analyse new algorithms and learning problems. In this scenario, however, it is difficult to see how this framework can be used to analyse the problem or new algorithms further. The framework used in [1] is simplified but is still within a setting where models can be trained and tested to evaluate how well the sample complexities reflect the true performance within the simplified setting. [1] Transformers as Algorithms: Generalization and Stability in In-context Learning [2] What Can Transformers Learn In-Context? A Case Study of Simple Function Classes [3] Transformers Learn In-Context by Gradient Descent. 2022 ------------------------------------------- Typos: L266: Section ?? L270: That it -> That is L305: Section ?? I have given a score of 6 for now, but I am open to changing my score based on the authors' responses and discussions with other reviewers.

Questions

NA

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA

Summary

The paper presents a theoretical framework for in-context learnability. The framework is grounded in the Probably Approximately Correct (PAC) learning theory and provides the first-ever finite sample complexity results for the in-context learning setup. The authors' approach involves a pretraining phase followed by an in-context learning phase where the training examples of the downstream task are concatenated in the input. The paper argues that, under specific conditions, latent tasks in the pretraining distribution can be effectively learned via in-context learning without changing the model's weights, even when the input significantly deviates from the pretraining distribution. This finding aligns with recent empirical results and suggests that in-context learning is more about task identification than learning.

Strengths

- The paper is well-written and clearly presented; - The paper tackles an important and underexplored theoretical aspect of in-context learning in LLMs. - It presents the first-of-its-kind PAC-based framework for in-context learnability, which could pave the way for further theoretical investigations. . The paper links its theoretical analysis to recent empirical findings, providing a sound validation for its results.

Weaknesses

- While the paper discusses a pretraining phase and an in-context learning phase, it doesn't clearly address how the transition between these phases is managed or optimized. Also, the instruction-tuning phase might be missing to ensure the success of in-context learning and mixture of tasks as defined in the paper; - Some empirical valuations including flipping labels could also be valuable to add [1, 2]; [1] Min, Sewon, et al. "Metaicl: Learning to learn in context." arXiv preprint arXiv:2110.15943 (2021). [2] Wei, Jerry, et al. "Symbol tuning improves in-context learning in language models." arXiv preprint arXiv:2305.08298 (2023).

Questions

- What if the downstream task is OOD from the pretraining data (which is common to see in terms of the generalization ability of instruction-tuned large language models), will the assumption, and conclusion still hold in Lemma 1?

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

- The PAC framework employed may not cover all aspects or scenarios of in-context learning with LLMs (order, variance in prompts, chain-of-thought prompting, generalization, emergent abilities, and so on). - Some inconsistencies in line 266, 305 with the appendix could be improved.

Summary

This paper studies the PAC-learnability of in-context learning when the pretraining distribution is a mixture of latent tasks, and the downstream task belongs to one of them. In addition to this mixure-of-tasks assumption, the other non-trivial assumptions the authors make include: the pretrained model can approximate the pretraining distribution arbitrarily close with a polynomial number of examples, and that the adjacent strings to be concatenated are approximately independent. Based on these assumptions, the author provides sample complexity results to demonstrate in-context learnability. Particularly, the correctness of the downstream task labels does not affect the validity of the conclusions.

Strengths

* The problem addressed in this paper, namely theoretical understanding of in-context learning, is undoubtedly important and timely. * The author addresses the PAC-learnability of in-context learning, which has not been well-studied. The problem formulation and results proposed by the authors are also interesting and novel. * The overall presentation is clear, with clear motivations and summaries for each part (although several details need to be improved, see Weaknesses).

Weaknesses

* The validity of the overarching mixture-of-tasks assumption for the pretraining distribution may need more justification. It appears to oversimplify the problem given the messy pretraining corpora and the novel downstream tasks used in practice (although it generalizes the similar assumption made by Xie et al. 2022, this limitation still exists). * The insensitivity to the correctness of in-context examples' labels may contradict recent observations. For instance, Wei et al. 2023 demonstrate that LLMs can do linear regression in in-context learning, which is clearly dependent on labels and thus cannot be directly explained by this paper. * There are some issues with the current presentation and writing that impede readability, such as: - The input and output of function $f_\theta$ should be made clearer. Based on the context, its input seems to be a string, but in equation (4) it has a conditional-style input which is not explained in the main text. - The meaning of $x$, $o$, and $s$ should be made clearer, preferably with some examples. The Kleene star symbol should also be clarified. - Typos, for instance run-away references to the appendix, Line 297 "other another", Line 362 "works designed" -> "works are designed".

Questions

Is the framework extendable to explain the observation where in-context learning can do label-sensitive linear regression?

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have not discussed the limitations.

Summary

This paper aims to explain the learnability of in-context learning. The main idea is that the pretraining tasks learn multiple downstream tasks and the prompt specify a particular task.

Strengths

1. The paper addresses the learnability of the in-context learning. It's a hot topic and very few theories are constructed from my knowledge. 2. The authors made several assumptions on in context learning and explain why they make the assumptions clearly.

Weaknesses

Though I am positive towards the paper, I think the main weakness is that the theoretical framework is too far away from the practice. The paper tells a good story but I am not sure what kind of guidance can the theory provide for the in-context learning in practice. For example, it looks that if we provide larger number of examples in prompts, the bound may be tighter. Is this really true for in-context learning? The authors can totally claim that there is computation limitation so we cannot make the number of examples to be as large as possible in practice. But if so, we also do not need in-context learning. The authors kind of explain the mystery of in-context learning. But kind of not. If we can make the number of examples infinity, supervised learning should also work well. I think the gap here is that the authors do not explain why this kind of in-context learning is better than supervised learning. Overall, I think the contribution is interesting but not significant. Minor: There are some ??'s in the paper. The authors probably want to fix them.

Questions

See the weakness section.

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

This is a theory paper. Societal impact is irrelevant.

Summary

In-Context Learning (ICL) allows large language models (LLMs) to be easily specialized to natural language downstream tasks. When users input a concatenated string of examples of a particular downstream task, modern LLMs often perform successfully without changing their weights, providing an effective new angle to tackle multiple NLP tasks immediately without fine-tuning. This paper defines ICL within the PAC learning framework. The authors find that pretrained LLMs learn a mixture distribution of downstream tasks (though these models maximize the likelihood of self-supervised next tokens or masked tokens on the training corpus). Under the mild assumptions, they show that ICL is provably guaranteed for LLMs to uncover the latent task, improving their performance without modifying any weights.

Strengths

The papers accept the crucial observations while providing theoretical guarantees. - Even if pre-training is assumed to be a process of learning mixture distribution of downstream tasks, it could never be equivalent to fine-tuning on a target task due to other irrelevant tasks. - Providing concatenation of independent examples is not natural for pretrained LLMs because they have never encountered such examples while pre-training.

Weaknesses

While some assumptions are claimed mild for learnability context, there are other assumptions which could not be easily acceptable given the practical behaviors of LLMs. - To make the provable bounds practically useful, both $c_1$ must be close to 1. However, the two strings delimited by the newline are not necessarily on paragraph levels. See the more details on the questions. - If we increase the size of vocabulary, it would be easier to guarantee the existence of positive $c_2$, but there would be a computational bottleneck while $c_2$ could be still very low, making the bound less practical. See the more details on the questions.

Questions

There are two questions mainly for the Assumption 2 and 3. (Regarding the Assumption 2) It is true that we can always find a positive c_1 that can easily bound the fraction in Equation (5) given $s_1$ and $s_2$. One concern is that users of LLM often use the newline character either trying to distinguish different shots of examples or differentiate consecutive paragraphs. - For the former case, would it be better to bound min and max of the fraction provided with two different orders $(s_1, \n, s_2)$ and $(s_2, \n, s_1)$? This is because the order of few-shot examples actually matter for the ICL performance. - For the latter case, it is difficult to say that two consecutive paragraphs would be relatively independent. It could be independent in a particular task, but not generally in most NLP tasks. For example, Coreference resolution is barely runnable without the previous paragraphs. In addition, it seems $c_1$ is defined as a strongly uniform constant (rather than depending on the choice of $s_1$ and $s_2$, and even including any possible NLP task). Would it be too strong? All in all, $c_1$ is likely to be extremely small number, thus making the bound less practical. (Regarding the Assumption 3) If the size of vocabulary is very large (typically 50k tokens), it becomes easier to find such $c_2$ (likely another uniform constant), but the $c_2$ would be very small yet requiring computational costs to handle large vocabulary. If the size of vocabulary is not large, it becomes harder to guarantee the existence of such $c_2$. Overall, $c_2$ would likely become a tiny positive number though the assumption would hold. (Regarding the Lemma 1) How to justify the last line of Lemma 1, where $m_{\hat{D}}$ can be chosen to be polynomial in both $\log \frac{1}{c_1 \cdot c_2}$ and $\frac{1}{\Delta_{KL}}$ within the context of Equation in Line 252? If so, would be the polynomial time both lower- and upper-bounded by the $\log \frac{1}{c_1 \cdot c_2}$? (Minor questions) - Section reference in Line 266 is broken.

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

No specific points are described or probed.

Thanks for your rebuttal.

Thanks the authors for their additional work to clarify my comments and feedback. Indeed this work will be better shine by addressing some of your extension and tighter bounds with respect to my suggestions. Hope these will be included in the final draft and future directions. Upon this feedback, I have increased my score.