SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Noam Wies, Yoav Levine & Amnon Shashua
{noam.wies,yoav.levine,shashua}@cs.huji.ac.il
In-context learning is a surprising and important phenomenon that emerged when modern language models were scaled to billions of learned parameters. Without modifying a large language model's weights, it can be tuned to perform various downstream natural language tasks simply by including concatenated training examples of these tasks in its input. Though disruptive for many practical applications of large language models, this emergent learning paradigm is not well understood from a theoretical perspective. In this paper, we propose a first-of-its-kind PAC based framework for in-context learnability, and use it to provide the first finite sample complexity results for the in-context learning setup. Our framework includes an initial pretraining phase, which fits a function to the pretraining distribution, and then a second in-context learning phase, which keeps this function constant and concatenates training examples of the downstream task in its input. We use our framework in order to prove that, under mild assumptions, when the pretraining distribution is a mixture of latent tasks (a model often considered for natural language pretraining), these tasks can be efficiently learned via in-context learning, even though the model's weights are unchanged and the input significantly diverges from the pretraining distribution. Our theoretical analysis reveals that in this setting, in-context learning is more about identifying the task than about learning it, a result which is in line with a series of recent empirical findings. We hope that the in-context learnability framework presented in this paper will facilitate future progress towards a deeper understanding of this important new learning paradigm.
The practice of pretraining language models (LMs) over massive general purpose text corpora has revolutionized natural language processing in recent years (Radford et al., 2019; Devlin et al., 2019; Brown et al., 2020). After an LM has been pretrained, the common approach for applying it to a specific downstream natural language task is to further train it on task-specific data, in a procedure referred to as fine-tuning (Howard & Ruder, 2018; Radford et al., 2018; Devlin et al., 2019). In their influential GPT3 paper, Brown et al. (2020) showed that a non-trivial alternative to fine-tuning emerges when the LM is large enough: an LM can be specialized to a downstream natural language task by simply receiving in its input a string composed of concatenated training examples of this task. Importantly, while the LM's weights are unchanged in this procedure, some form of learning evidently takes place; the performance on the downstream task was shown to significantly improve with the number of concatenated training examples, for a disparate variety of natural language tasks.
This phenomenon, referred to as in-context learning, has had a profound practical impact on the applicability of large LMs: one does not need to have any access to the model weights in order to specialize the model for a certain task. Instead, a string of training examples provided even via API access to the model is enough, and often not many examples are required (Brown et al., 2020). However, despite its growing popularity in a multitude of use-cases (Bommasani et al., 2021), the reasons for the effectiveness of in-context learning in pretrained LMs are not well understood from
a theoretical perspective. In particular, this new learning paradigm still lacks a formal definition of learning.
In this paper, we propose a PAC learning (Valiant, 1984) definition for in-context learning, along with the first finite sample-complexity learning results for in-context learning. Our results shed light on the mysterious question of why in-context learning works even though a string of concatenated input-output pairs does not resemble the natural distribution of pretraining examples. Our framework is based on an interpretation of pretraining as unsupervised multi-task learning of many natural language tasks, which dates back to the GPT2 paper (Radford et al., 2019), one of the instigators of the LM pretraining era. This view, which has since been widely adopted, was supported by experiments that revealed non-trivial zero-shot capabilities of pretrained LMs on a variety of natural language tasks. In other words, an LM pretrained only to maximize the likelihood of unlabeled text from the web, was able to perform reasonably well on a variety of downstream natural language tasks without learning from any additional training examples post pretraining, implying that many different natural tasks are directly learned during pretraining.
Following this interpretation, we prove our in-context learning results in a setting in which the pretraining distribution is a mixture of downstream tasks. Importantly, the pretraining examples are not explicitly associated with downstream tasks, but rather they are drawn from a mixture of tasks and the association between examples and tasks is latent. We prove that under mild assumptions, in-context learning is guaranteed to happen for a model trained by such multi-task pretraining. We show that the in-context learning mechanism of concatenating input/output pairs of a certain task allows the pretrained LM to uncover the latent task and improve its performance without modifying its weights. Our assumptions are quite general, namely that there is a lower bound on the probability of any mixture component and of any token, and that the downstream tasks are both distinguishable and have sufficient margin between different labels. See Section 3.1 for a formal definition of our assumptions.
The interpretation of in-context learning as uncovering tasks that were already learned during pretraining is also supported by empirical evidence. For example, both Min et al. (2022) and Lyu et al. (2023) showed that replacing the labels provided in the in-context training examples with random task labels barely affects the performance of in-context learning, implying that the in-context learning mechanism is more about identifying the task than about learning it. Similarly, both Webson & Pavlick (2022) and Lampinen et al. (2022) studied whether LMs truly understand the text of the in-context examples, and found that irrelevant text that mimic the task can have a similar effect as learning with true training examples.
Finally, the latent task inference perspective of in-context learning was also examined theoretically. Xie et al. (2022) study a language-inspired toy distribution for pretraining. Specifically, they analyze a pretraining distribution consisting of a mixture of Hidden Markov Models (HMMs), where each HMM factor corresponds to a single task. They showed that in this setting in-context learning was guaranteed, however, unlike our work, their analysis is confined only to the infinite limit of the number of in-context examples. We provide polynomial sample complexity guarantees for in-context learning, which are much more relevant to its efficiency in practice (often referred to as "few-shot"). Furthermore, Xie et al. (2022) limit their analysis only to a certain class of mixture of HMMs as pretraining distribution, while our results are for any mixture distribution under mild assumptions (see above). Lastly, their analysis assumes perfect learning of the pretraining distribution while our PAC approach captures imperfect pretraining. Beyond Xie et al. (2022), several recent papers (Akyurek et al., 2023; Dai et al., 2022; von Oswald et al., 2022) showed that from an expressivity point of view, self-attention architectures can implement a gradient descent algorithm with in-context examples. However, these works do not justify why pretraining on natural data will converge to weights that implement gradient decent, nor do they provide finite sample-complexity results for in-context learning. Overall, our presented PAC framework for in-context learning (Section 2) allows us to prove the first finite in-context learning sample-complexity results, in a quite general setting (Section 3). We hope that this framework will facilitate a further expansions of the understanding of in-context learning as a new learning paradigm.
2 A PAC Learnability Framework for In-Context Learning
In this section, we define a Probably Approximately Correct (PAC) learnability framework (Valiant, 1984) for in-context learning. Conceptually, we aim to adjust the PAC learnability framework, in order to capture the in-context few-shot learning capabilities of large language models (Brown et al., 2020; Wei et al., 2022a; Sanh et al., 2022; Ouyang et al., 2022).
We begin by describing the in-context learning paradigm. Let  $f_{\theta}$  be a function fitted to a massive general purpose pretraining distribution. In order to apply  $f_{\theta}$  for a specific downstream task, it is common to further train  $f_{\theta}$  on pairs of the task's inputs  $x_{1}, \ldots, x_{k}$  with their corresponding labels  $y_{1}, \ldots, y_{k}$ , in a procedure referred to as fine-tuning. In-context learning is an alternative, simpler technique, which "teaches"  $f_{\theta}$  the task at hand by constructing a prompt comprised of concatenated training examples. Specifically, denote the string concatenation operator by  $\oplus$ , and let " $\backslash n$  " be a special delimiter token. Then, the in-context learning technique is to construct a prompt  $p$  by concatenating pairs of the task's inputs along with their labels:
$$
p := x _ {1} \oplus y _ {1} \oplus “ \backslash n ” \oplus \dots \oplus x _ {k} \oplus y _ {k} \oplus “ \backslash n ” \tag {1}
$$
Given this prompt, the prediction is made by choosing the label  $y$  that is the most likely continuation for the prefix  $p \oplus x$  according to the function  $f_{\theta}$ . In other words, by predicting the label  $y$  that maximizes  $f_{\theta}(p \oplus x \oplus y)$ . Notably, the prefix  $p \oplus x$  does not resemble inputs that  $f_{\theta}$  has trained on. Note that we used the newline symbol "\n" for clarity. However, in practice the delimiter token might depend on implementation and could, for instance, be two successive newlines "\n\nas done in the few-shot learning evaluation framework Evaluation Harness (Gao et al., 2021).
We now present our in-context learning learnability definition. Let  $\mathcal{D}$  be a pretraining distribution over strings of characters from an alphabet  $\Sigma$ . And let  $\tilde{\mathcal{D}}$  be a downstream task's distribution over pairs of strings  $x \in \Sigma^{\star}$  and theirs corresponding labels  $y \in \Sigma^{\star}$ . We aim to define distribution-dependent learnability in the in-context learning setup, for a frozen probabilistic model  $f_{\theta}$  that originally trained to maximize the likelihood of the pretraining distribution  $\mathcal{D}$ . Then, given a prompt  $p \sim \tilde{\mathcal{D}}^{k}$  which consist of  $k$  pairs of inputs and theirs corresponding labels, each independently sampled from  $\tilde{\mathcal{D}}$ , the model  $f_{\theta}$  is tested on the zero-one loss of the in-context predictor:
$$
L _ {\text {i n - c o n t e x t}, \tilde {D}} := \mathbb {E} _ {x, y \sim \tilde {\mathcal {D}}} \left[ l _ {0 - 1} \left(\operatorname {a r g m a x} _ {y ^ {\prime}} f _ {\theta} \left(p \oplus x \oplus y ^ {\prime}\right), y\right) \right] \tag {2}
$$
Since we are interested in frozen models as opposed to models that are fine-tuned, we will require that the same  $f_{\theta}$  model achieves a low  $L_{\mathrm{in-context},\tilde{\mathcal{D}}}$  loss for multiple downstream tasks simultaneously. The following is our PAC learning definition for in-context learning of a model that was pretrained to maximize the likelihood of the pretraining distribution.
Definition 1. Let  $\mathcal{H}$  and  $\tilde{\mathcal{H}}$  be hypothesis classes of pretraining distributions and downstream distributions respectively. We say that  $\tilde{\mathcal{H}}$  is in-context learnable after pretraining on a pretraining distribution  $\mathcal{D} \in \mathcal{H}$ , if there exist a pretraining algorithm  $\mathcal{A}$  and a pair of functions  $m_{\mathcal{H}}, m_{\tilde{\mathcal{H}}} : (0,1)^2 \to \mathbb{N}$  with the following property: For every  $\epsilon, \delta > 0$  and every  $\tilde{\mathcal{D}} \in \tilde{\mathcal{H}}$ , if the number of pretraining examples  $n$  is greater than  $m_{\mathcal{H}}(\epsilon, \delta)$ , and if the number of the downstream tasks' examples  $k$  is greater than  $m_{\tilde{\mathcal{H}}}(\epsilon, \delta)$ , then with probability of at least  $1 - \delta$  (over the choice of the  $n + k$  examples) the following holds:
$$
L _ {i n - c o n t e x t, \tilde {\mathcal {D}}} - B a y e s E r r o r R a t e \leq \epsilon \tag {3}
$$
In addition, we will say that a collection of downstream tasks  $\tilde{\mathcal{H}}$  is efficiently in-context learnable after pretraining on  $\mathcal{D}$ , if both  $m_{\mathcal{H}}$  and  $m_{\tilde{\mathcal{H}}}$  in the above definition are polynomial in  $\log \epsilon^{-1}$ ,  $\log \delta^{-1}$  and  $\log |\tilde{\mathcal{H}}|$ .
In order to focus on in-context learning questions that are relevant for the natural language processing setup, and avoid complications that are not necessary for answering them, we will assume that the pretraining distribution  $\mathcal{D}$  belongs to some hypothesis class  $\mathcal{H}$  that is learnable by some pretraining algorithm:
Assumption 1. There exists a learning algorithm  $\mathcal{A}$ , and a sample complexity function  $m_{\mathcal{D}}: (0,1)^2 \to \mathbb{N}$ , with the following property: for any pretraining distribution  $\mathcal{D} \in \mathcal{H}$  and any  $\epsilon, \delta > 0$ , if the number of pretraining examples  $n$  is greater than  $m_{\mathcal{D}}(\epsilon, \delta)$ , then with probability at least  $1 - \delta$
over the pretraining examples, for any  $T \geq 1$ , the total variation of the conditional distributions of the  $T$  'th token is at most  $\epsilon$ :
$$
\max  _ {o _ {1} \dots o _ {T} \in \Sigma} \left| \mathbb {P} _ {\mathcal {D}} \left(o _ {T} \mid o _ {1} \dots o _ {T - 1}\right) - f _ {\theta} \left(o _ {T} \mid o _ {1} \dots o _ {T - 1}\right) \right| <   \epsilon \tag {4}
$$
where  $f_{\theta}$  denotes the model yielded by the pretraining algorithm  $\mathcal{A}$  and  $f_{\theta}(o_{T}|o_{1}\dots o_{T - 1})$  denotes for the conditional likelihood  $\frac{f_{\theta}(o_1\dots o_{T - 1}o_T)}{f_{\theta}(o_1\dots o_{T - 1})}$ . In addition,  $m_{\mathcal{D}}$  is polynomial in both  $\epsilon^{-1}$  and  $\delta^{-1}$ .
With the above formal definition of in-context learning, we aim to shed some light on the mysterious in-context learning abilities of large LMs. Specifically, in the next section, we will shed light on the following major question: How do frozen pretrained models learn from prompts that do not resemble their pretraining distribution?
3 Guarantees on In-Context Learning
In this section, we will demonstrate the use of the above learnability framework for in-context learning and analyze a setting in which pretraining a simple language model provably leads to in-context learning capabilities. For doing so, we follow previous work and view the language modeling task as an implicit multi-task setup. Indeed, creating human-like text involves many different skills, from grammar to world knowledge, so learning a language model inevitably develops a variety of skills (Gulordava et al., 2018; Zhang & Bowman, 2018; Weber et al., 2021). The implicit unsupervised multi-task view of language modeling can be traced to the GPT2 paper (Radford et al., 2019), which revealed that pretrained LMs are capable of a wide variety of natural language tasks without the need for further training. Since then, this view has been reinforced for example by Hendrycks et al. (2021). They showed that on a diverse massive set of 57 real world text understanding tasks, the largest GPT-3 (Brown et al., 2020) language model improves over random chance by almost 13 percent points on average. Importantly, these results were obtained in a zero-shot learning setting, i.e. in a setting where the language model had only been pretrained to maximize the likelihood of unlabeled text. Accordingly, these results suggest that many different natural tasks are directly learned during pretraining.
We reflect the above implicit multi-task nature of language modeling by assuming that the pretraining distribution contains a latent variable that represents the task at hand. We show below that for such pretraining distributions, adding training examples to the in-context learning prompt implicitly reveals the already learned latent task. In Subsection 3.1, we present a multi-task pretraining hypothesis class, for which Subsection 3.2 shows that in-context learning reveals what task is currently being performed.
3.1 The analyzed latent concept hypothesis class
In this subsection, we describe the analyzed multi-task pretraining hypothesis class, as well as the corresponding downstream task hypothesis class that can be learned in-context after pretraining. To begin, we define the pretraining distribution as a mixture of multiple downstream tasks. Importantly, during pretraining the downstream task of each example is unknown, i.e. it is a latent variable, and thus pretraining is not equivalent to fine-tuning on the task since the model cannot simply ignore pretraining examples of irrelevant tasks. Specifically, we generate a length- $T$  pretraining example  $x_{1},\ldots ,x_{T}\in \Sigma$  from the pretraining distribution  $\mathcal{D}$  by first sampling a concept  $\phi$  from a family of concepts  $\Phi$  according to a prior  $\mathbb{P}(\phi)$ . We then sample the tokens according to the concept's specific distribution  $\mathbb{P}_{\phi}(x_1,\dots,x_T)$ .
Moving to the downstream tasks, we will prove in-context learnability results for tasks where the underlying inputs distribution is a component  $\phi$  in the pretraining mixture distribution  $\mathcal{D}$ . Formally, we generate a length  $T$  downstream task example  $x$  and the corresponding label  $y$  from  $\tilde{\mathcal{D}}$  by first sampling  $T$  tokens  $o_1, \ldots, o_T$  according to the downstream task distribution  $\mathbb{P}_{\phi}(o_1, \ldots, o_T)$ . Then we assemble  $x$  using all tokens except the last one, and set the label to be that token i.e. we set that  $x_t = o_t$  for any  $t < T$  and that  $y = o_T$ . Note that in principle the concatenation of independent examples in  $p$  causes a distribution drift from the pretraining distribution<sup>1</sup>. In this sense, the analyzed
Formally, the distribution drift is caused by the fact that the concatenation of the " $\backslash n$ " after  $y$  ignores the marginal probabilities of " $\backslash n$ ".
model captures the fact that few-shot prompts are unnatural since they are not encountered during pretraining.
Now we describe assumptions about the pretraining distributions, for which we will prove our in-context learnability results. Our first assumption requires that given a delimiter token "\n", two successive strings  $s_1$  and  $s_2$  that are concatenated with "\n" are approximately independent according to  $\mathbb{P}_{\phi}$ :
Assumption 2. There exists a constant  $0 < c_{1} \leq 1$  such that for any two strings  $s_{1}, s_{2}$  in  $\Sigma^{\star}$  and any concept  $\phi$  the following holds:
$$
c _ {1} \leq \frac {\mathbb {P} _ {\phi} \left(s _ {1} \oplus {} ^ {\prime \prime} \backslash n ^ {\prime \prime}\right) \cdot \mathbb {P} _ {\phi} \left(s _ {2}\right)}{\mathbb {P} _ {\phi} \left(s _ {1} \oplus {} ^ {\prime \prime} \backslash n ^ {\prime \prime} \oplus s _ {2}\right)} \leq \frac {1}{c _ {1}} \tag {5}
$$
Note that when the two successive strings  $s_1$  and  $s_2$  are exactly independent, the probability ratio in Equation 5 is equal to one, and the constant  $c_1$  quantifies the deviation from this situation. While this assumption might sound restrictive, it is reasonable to assume that two consecutive paragraphs are not highly dependent according to the distributions in the pretraining mixture. Intuitively, we will use this assumption in order to apply concentration inequalities to the likelihood of the in-context prompt, and we will deduce that the role of the in-context prompt is to reweight the prior regarding the different mixture components. It is important to note that the approximate independence assumption for any component in the mixture does not imply approximate independence of the mixture distribution itself. Hence this reweighting is possible, since the assumption does not imply that the in-context prompt is ignored.
Beyond this approximate independence, we will also require that there exist a lower bound on the conditional probability of any single token:
Assumption 3. There exist a constant  $c_2 > 0$  such that for any string  $s$  in  $\Sigma^{\star}$ , any character  $\sigma \in \Sigma$  and any concept  $\phi$  the following holds:
$$
\mathbb {P} _ {\phi} (\sigma \mid s) > c _ {2} \tag {6}
$$
Basically, we need this assumption in order to avoid the harm of zero likelihood due to the unnatural concatenation of input and output pairs in the prompt  $p$  (where the prompt  $p$  is defined in Section 2). Note that without such an assumption, in-context learning is impossible since the probability of a prompt  $p$  might be zero and hence the prediction of the model in such cases becomes meaningless. Finally, we assume that the prior distribution is strictly positive. In other words, we will assume that there is a lower bound higher than zero on the likelihood of any concept appearing in the pretraining distribution.
Assumption 4. There exist a constant  $c_3 > 0$  such that for the prior  $\mathbb{P}_{\mathcal{D}}(\phi)$  of any concept  $\phi$ , is at least  $c_3$ .
Clearly, without such an assumption, the prior of the downstream task can be arbitrarily low, which means it will be nearly impossible to recognize the task. In the next subsection, we will use the above assumptions in order to provide in-context learning guarantees via the mechanism of task recognition.
3.2 Guarantees on In-Context Learning via Latent Concept Inference
In this subsection, we analyze the prediction of an in-context learning model in the setting described in Subsection 3.1. We show that in this setting, there is a polynomial sample complexity that guarantees in-context learning is Probably Approximately Correct. At a high level, since the pretraining is not precise and can only approximate the pretraining distribution  $\mathcal{D}$  up to some error  $\triangle_{\mathrm{pretraining}} > 0$  (see Assumption 1), we will split the in-context prediction analysis into two parts. The first part will involve the simpler case of test examples  $x$  and corresponding label candidates  $y$  and  $\tilde{y}$  for which the margin between the conditional likelihoods of the labels  $\mathbb{P}_{\tilde{\mathcal{D}}}(\boldsymbol {y}|\boldsymbol {x})$  and  $\mathbb{P}_{\tilde{\mathcal{D}}}(\tilde{y} |x)$  is large enough. In this scenario, we show that both the deviation due to imperfect pretraining and due to imperfect task recognition is negligible. Therefore, we will conclude that such deviations do not have any impact on the loss of in-context learning. In the second scenario, when the difference between ground-truth likelihoods is small, the error rate of the Bayes optimal classifier must be high. Accordingly, even though the in-context predictor might confuse between labels, the loss in any case will be small because we only compare it to the error rate of the Bayes optimal classifier.