SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Starting with the first scenario, where the margin between label candidates is sufficiently large, we will prove that as more examples are added to the prompt, the in-context predictions converge to the correct label. As a preliminary step, we prove a lemma regarding the ratio of the prompt likelihoods  $\mathbb{P}_{\phi}(p)$  (where the prompt  $p$  is defined in Section 2) according to the ground-truth task versus the other mixture components. In other words, we ask how likely it is that the prompt of concatenated examples was sampled according to one of the tasks distributions versus the other tasks distributions. We will use this lemma in order to estimate the effect of the in-context prompt, on the prior regarding the different mixture components. Specifically, we denote by  $\triangle_{\mathrm{KL}}$  the minimum Kullback-Leibler divergence between the ground-truth component, and the other mixture components. Then, we prove that the ratio of the prompt probabilities according to the ground-truth components converge to zero, with rate that is exponential in both the number of in-context examples  $k$ , and in the minimal Kullback-Leibler divergence  $\triangle_{\mathrm{KL}}$ . Intuitively, the exponential rate with regard to the number of examples comes naturally from the fact that each example in the prompt is sampled independently of the others. As a result, their effect is a multiplicative one. Additionally, the Kullback-Leibler divergence between the ground truth component and the other mixture components measures log probabilities, while we are interested in the probabilities themselves. Hence the rate is also exponential in the above Kullback-Leibler divergence. Formally, we have that:
Lemma 1. Let  $\mathcal{D}$  be a pretraining distribution for which assumptions 2,3 hold, and let  $\phi^{\star}$  be a downstream task mixture component from  $\mathcal{D}$  such that  $\triangle_{KL} > 8\log \frac{1}{c_1\cdot c_2}$ . Then, there exists  $m_{\tilde{\mathcal{D}}}: (0,1)^2 \to \mathbb{N}$  with the following property: for any  $\epsilon, \delta > 0$  and any  $\phi \neq \phi^{\star}$ , if the number of in-context examples  $k$  is at least  $m_{\tilde{\mathcal{D}}}(\epsilon, \delta)$ , then,  $\frac{\mathbb{P}_{\phi}(p)}{\mathbb{P}_{\phi^{\star}}(p)} < \epsilon$  with probability of at least  $1 - \delta$  (over the choice of the  $k$  in-context examples). Moreover, the above still holds when the labels in  $p$  are randomly flipped, and  $m_{\tilde{\mathcal{D}}}$  can be chosen such that it will be polynomial in  $\log \frac{1}{\delta}$ ,  $\log \frac{1}{\epsilon}$ ,  $\log \frac{1}{c_1\cdot c_2}$ ,  $\frac{1}{\triangle_{KL}}$  and  $T$ .
Proof. In essence, we prove this lemma by using the approximate independence assumption in order to apply concentration inequalities. In addition, we use Assumption 3 for bounding the distribution drift that is caused by the artificially inserted newline token. In particular, the log of the ratio of prompt probabilities according to different mixture components is concentrated around its expectation. Importantly, this expectation is equal to minus one times the Kullback-Leibler divergence between the components, plus a term that is caused by the mentioned distribution drift. Thus, we conclude that the ratio of the prompt probabilities according to the ground-truth task distribution and the other pretraining components converges to zero. Moreover, the rate of that convergence is exponential in both the number of in-context demonstrations, and the minimal Kullback-Leibler divergence between different mixture components. See full details in Section A of the appendix.
Now we will use the above Lemma 1 to analyze the in-context predictions, namely the labels  $\tilde{y}$  that maximize the likelihood of the concatenation of the in-context prompt  $p$ , with the example  $x \oplus \tilde{y}$  according to the pretrained model  $f_{\theta}$ . Essentially, we will aim to understand when these predictions are identical to the Bayes Optimal Classifier predictions. That it, when these predictions align with the labels  $y$  that maximize the likelihood of the example  $x \oplus y$  as determined by downstream task distribution  $\mathbb{P}_{\tilde{\mathcal{D}}}$ . Since we are in the first scenario, where the margin between label candidates is large enough. We will prove that in this case, for large enough  $k$ , the ground truth in-context predictor also has margin that is at least half of the original margin. Hence, we will conclude that for such  $k$  the ground truth in-context predictor is equal to the Bayes Optimal Classifier. Moreover, we will prove a lower bound on the margin of the in-context predictions, so the above still holds for any distribution that approximates the pretraining distribution sufficiently well.
Formally, for a test example input  $x$ , we define the margin  $\triangle (x,y,\tilde{y})$  between two label candidates  $y$  and  $\tilde{y}$  as the difference between their conditional likelihoods  $\mathbb{P}_{\tilde{\mathcal{D}}} (y|x)$  and  $\mathbb{P}_{\tilde{\mathcal{D}}} (\tilde{y} |x)$  according to the downstream ground-truth distribution. Similarly, for a test example input  $x$  and in-context prompt  $p$ , we define the margin  $\triangle (p,x,y,\tilde{y})$  between two label candidates  $y$  and  $\tilde{y}$  as the difference between their conditional likelihoods  $\mathbb{P}_{\mathcal{D}} (y|p\oplus x)$  and  $\mathbb{P}_{\mathcal{D}} (\tilde{y}|p\oplus x)$  according to the pretraining distribution conditioned on the prompt  $p$ . Using these definitions, we consider triplets of test example input  $x$  and two labels candidates  $y,\tilde{y}$  with margin at least two times the pretraining error  $\triangle_{\mathrm{pretraining}}$  as the first scenario. In this case we will prove that for large enough  $k$  the ground truth in-context predictor also has margin of at least the pretraining error  $\triangle_{\mathrm{pretraining}}$ . This means that deviations arising from imperfect task recognition will be negligible in this situation. Moreover, the margin will remain larger than the pretraining error, which means the pretrained model  $f_{\theta}$  handles this case well.
Theorem 1. Let  $\mathcal{D}$  and  $\tilde{\mathcal{D}}$  be a pair of pretraining distribution and downstream task, for which Assumption 4 as well as the assumptions in Lemma 1 upholds. Then, there exists  $m_{\tilde{\mathcal{D}}}:(0,1)^2\to \mathbb{N}$  with the following property: for every test example  $x$  and two label candidates  $y,\tilde{y}$  with positive margin  $\triangle (x,y,\tilde{y}) > 0$  and  $\delta >0$ , if the number of in-context examples  $k$  is at least  $m_{\tilde{\mathcal{D}}}\left(\frac{\triangle(x,y,\tilde{y})}{2},\delta\right)$ , then  $\triangle (p,x,y,\tilde{y}) > \frac{\triangle(x,y,\tilde{y})}{2} +c_1^2 -1$  with probability of at least  $1 - \delta$  (over the choice of the  $k$  in-context examples). Moreover, the above still holds when the labels in  $p$  are randomly flipped<sup>2</sup>, and  $m_{\tilde{\mathcal{D}}}$  can be chosen such that it will be polynomial in  $\log \frac{1}{\delta}$ ,  $\log \frac{1}{\epsilon}$ ,  $\log \frac{1}{c_1\cdot c_2\cdot c_3}$ ,  $\frac{1}{\Delta_{KL}}$  and  $T$ .
Proof. We begin by writing the difference between the ground-truth label likelihood, and the likelihood of other another label  $\tilde{y}$  explicitly. Specifically, by the definition of conditional probabilities we have that:
$$
\mathbb {P} _ {\mathcal {D}} (y \mid p \oplus x) - \mathbb {P} _ {\mathcal {D}} (\tilde {y} \mid p \oplus x) = \frac {\sum_ {\phi} \mathbb {P} _ {\mathcal {D}} (\phi) [ \mathbb {P} _ {\phi} (p \oplus x \oplus y) - \mathbb {P} _ {\phi} (p \oplus x \oplus \tilde {y}) ]}{\sum_ {\phi} \mathbb {P} _ {\mathcal {D}} (\phi) \mathbb {P} _ {\phi} (p \oplus x)} \tag {7}
$$
Now, denote by  $\phi^{\star}$  the mixture component of  $\tilde{D}$ , then Assumption 2 assures us that for each component in the mixture, the textual prompt is approximately independent of the test example. So we can use Lemma 1 and get that in both the numerator and the denominator, the  $\phi^{\star}$  term is the dominant term in the sum. Thus, we will get that the difference of the labels' likelihoods is approximately the fraction of the  $\phi^{\star}$  terms, which is equal to the original downstream task margin. So we will conclude that the margin of the in-context predictor is at least half of the downstream task original margin. See full details in Section B of the appendix.
We denote by  $\triangle_{\tilde{\mathcal{D}}}$  the minimal margin of the Bayes Optimal Classifier predictions in downstream task  $\tilde{\mathcal{D}}$  i.e. the minimal margin between the Bayes Optimal Classifier prediction and another labels. With the above definitions and theorem, we can combine the scenario of large margins, which are preserved by the in-context predictor, with the scenario of small margins, in which the loss is minimally affected by the wrong prediction, and prove our main in-context learnability results:
Theorem 2. Let  $\tilde{\mathcal{H}}$  be hypothesis classes of downstream distributions, and denote by  $\mathcal{D}$  a mixture distribution on  $\tilde{\mathcal{H}}$  for which assumptions 1,2,3 and 4 uphold. Further assume that the margin  $\triangle_{\tilde{\mathcal{D}}}$  of any downstream task  $\tilde{\mathcal{D}} \in \tilde{\mathcal{H}}$  is at least  $4 \cdot (1 - c_1^2)$ , and the minimal Kullback-Leibler divergence between different distributions is greater than  $\log \frac{1}{c_1 \cdot c_2}$ . Then  $\tilde{\mathcal{H}}$  is efficiently in-context learnable after pretraining on  $\mathcal{D}$  (see Definition 1).
Proof. Assumption 1 assures us the existence of the pretraining algorithm with a polynomial sample complexity. So we will prove the theorem with pretraining sample complexity that is derived from this algorithm, where the accuracy required from pretraining is  $\frac{2}{a \cdot (1 - \alpha)}$  times the accuracy required from in-context learning (see analysis below). In addition, Theorem 1 assures us that for large enough  $k$  the ground truth in-context predictions are equal to the Bayes Optimal Classifier predictions, so we will prove the theorem with downstream sample complexity that is derived from Theorem 1.
Let  $\epsilon, \delta > 0$  and denote by  $f_{\theta}$  the model that was learned during the pretraining process. We will prove that the contribution of any  $x$  to the loss  $L_{\mathrm{in-context},\tilde{D}}$  (see Equation 2) is at most  $\epsilon$  and hence complete the proof. Given  $x$ , let  $y$  be the Bayes Optimal Classifier prediction for that  $x$ . Then, we will split the analysis into two cases.
The first case is of examples  $x, y$  such that their margin  $\triangle (x, y, \tilde{y})$  from any alternative label candidate  $\tilde{y}$  is at least  $8 \cdot \triangle_{\text{pretraining}}$ . In this case, Theorem 1 assures us that for large enough  $k$  the ground truth in-context predictor also has margin  $\triangle (p, x, y, \tilde{y})$  that is greater than  $\frac{1}{2} \cdot \triangle (x, y, \tilde{y}) + c_1^2 - 1$ . Now since  $\triangle (x, y, \tilde{y}) \geq \triangle_{\hat{\mathcal{D}}} > 4 \cdot (1 - c_1^2)$  we have that  $\triangle (p, x, y, \tilde{y})$  is at least  $2 \cdot \triangle_{\text{pretraining}}$ . Thus, we conclude that in this case, the predictions of the ground truth in-context prediction and  $f_\theta$  are the same. Furthermore both of them are identical to the Bayes Optimal Classifier prediction, and hence  $x$  does not contribute to the loss.
2While theorem 1 provides label-insensitive results, it also predicts sensitivity to the labels in cases where the labels are the primary causes for the distinction between different tasks, as measured by the Kullback-Leibler divergence.
Moving to the second case, and denote by  $\tilde{y}$  the in-context prediction of  $f_{\theta}$ . Then, the arguments from the previous paragraph assure us that  $\mathbb{P}_{\phi^{\star}}(y|x) - \mathbb{P}_{\phi^{\star}}(\tilde{y} |x) < 8 \cdot \triangle_{\mathrm{pretraining}}$ , since otherwise we will get that  $f_{\theta}(y|x) > f_{\theta}(\tilde{y}|x)$ . Finally, we will choose that  $\triangle_{\mathrm{pretraining}} = \frac{\epsilon}{8}$ , and hence get that also in the second case, the contribution of  $x$  to the loss is less than  $\epsilon$ .
Since the Probably Approximately Correct (PAC) learning framework was introduced in Valiant (1984), a rich line of works has extended the framework to distribution-dependent bounds (Benedek & Itai, 1991; Vayatis & Azencott, 1999; Sabato et al., 2013) inter alia. These works relax PAC's adversarial requirement of generalization for all input distributions, and only consider distributions that satisfy certain statistical properties. Consequently, these papers provide more realistic sample complexity bounds. In this paper we present a framework for in-context learning learnability that uses self-supervised pretraining to assist in solving downstream tasks. In this framework, we follow the above line of distribution-specific works with the distinct feature of a self-supervised pretraining phase, and of learning with frozen models through their input context. The advantages of such pretraining have been studied from a theoretical perspective before. For example, Saunshi et al. (2021) and Wei et al. (2021) demonstrate that language modeling can benefit downstream tasks either by prompt tuning or head tuning. However, unlike our work, in their analysis some weights are learned. In contrast, we focus on frozen models that can learn only from their input context, which includes the concatenation of few inputs and outputs pairs.
Returning to in-context learning, a recent line of work study learning paradigms that are based on concatenating several inputs into a single input context. For example, Garg et al. (2022) show that a Transformer architecture is able to discover in-context learning algorithms for simple function classes, such as two-layer neural networks, and decision trees. Similarly, Akyurek et al. (2023) show that a Transformer architecture is able to discover efficient Bayes optimal least square learning algorithms, and Laskin et al. (2022) show that a Transformer architecture is able to discover efficient in-context reinforcement learning algorithms. Additionally, Levine et al. (2022) studied the inductive bias of in-context learning, and proved that Transformer based language models can model much stronger dependencies between text segments that appeared in the same training example. Finally, Li et al. (2023) proved a multi-task generalization bound for in-context learning with the Transformer architecture. However, unlike our work, the pretraining distribution in all of the above works designed to include input context that includes few-shot demonstrations explicitly. As a consequence, these methods do not answer the mysterious question of how frozen pretrained models can learn from in-context prompts that do not resemble their pretraining distribution.
Another recent line of work investigates possible mechanisms that enable in-context learning. For example, Olsson et al. (2022) provide indirect evidence that induction heads might constitute the mechanism for in-context learning in large transformer models. In addition, 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. Finally, Chan et al. (2022) provided empirical evidence that some data distributional properties encourage in-context learning even when the frozen models do not see prompts during pretraining. However, unlike our work, all of the above works does not provide theoretical guarantees of in-context learnability.
Finally, Xie et al. (2022) are the closest to our work. They studied a language-inspired toy distribution for pretraining and analyzed a pretraining distribution consisting of a mixture of Hidden Markov Models (HMMs). Unlike their results, our results are applicable to any mixture distribution that meets our mild assumption. In addition, unlike our polynomial sample complexity guarantees, their analysis guarantees in-context learning only for an infinite number of in-context examples. Lastly, their analysis assumes perfect pretraining distribution learning, but our approach captures imperfect pretraining.
The discovery of in-context learning in large LMs, made by Brown et al. (2020), was surprising to many in our field. A model that was pretrained to maximize the likelihood of natural text was able
to make use of concatenated training examples of downstream natural language tasks—inputs that do not resemble its pretraining distribution, and moreover these inputs improved the model's ability to perform the task. Our theoretical results, based on a common latent multitask framework for the pretraining phase, shed light on the above surprising mysteries. With our PAC-based framework, we were able to provide sample complexity guarantees for in-context learning in such pretrained models, which are not only the first finite sample complexity results for this framework but they also indicate efficient (polynomial) in-context learning, which reflect the behavior of this setting in practice.
We hope that our framework can be used to deepen the understanding of the in-context learning phenomenon. In particular, we mark the connection between model size and the in-context learning efficiency as an interesting open question (Wei et al., 2022b). Additionally, in-context learning has shown to be capable of learning new tasks not included in the pre-training distribution (Wei et al., 2023). The extensions of our results to such situations, as well as input-dependent bounds that capture the sensitivity of few-shot in-context learning to the order of the few-shot examples is an interesting open questions.
Limitations: While we assume that the pre-training distribution perfectly matches a mixture of downstream tasks, we acknowledge that the pre-training data in real-world LLMs is often noisy and imperfect. Our results represent an idealized scenario with perfect pre-training data. Extending our analysis to account for limitations of real-world pre-training data remains an important direction for future work.
The authors would like to thank Yotam Wolf and Oshri Avnery for their assistance and advice. This research was supported by the ERC (European Research Council) and the ISF (Israel Science Foundation).
Ekin Akyurek, Jacob Andreas, Dale Schuurmans, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=0g0X4H8yN4I.
Gyora M Benedek and Alon Itai. Learnability with respect to fixed distributions. Theoretical Computer Science, 86(2):377-389, 1991.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/1457c0d6bcbd4967418bfb8ac142f64a-Paper.pdf.
Stephanie C.Y. Chan, Adam Santoro, Andrew Kyle Lampinen, Jane X Wang, Aaditya K Singh, Pierre Harvey Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=1Hj-q9BSRjF.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559, 2022.
Colin Wei, Sang Michael Xie, and Tengyu Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 16158-16170. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/file/86b3e165b8154656a71ffe8a327ded7d-Paper.pdf.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022a. URL https://openreview.net/forum?id=gEZrGCozdqR.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022b. ISSN 2835-8856. URL https://openreview.net/forum?id=yzkSU5zdwD. Survey Certification.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=RdJVFCHjUMI.
Kelly Zhang and Samuel Bowman. Language modeling teaches you more than translation does: Lessons learned through auxiliary syntactic task analysis. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 359-361, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.18653/v1/W18-5448. URL https://aclanthology.org/W18-5448.
We begin by explicitly writing  $\mathbb{P}_{\phi}(p)$ . Specifically, Assumption 2 assures us that each in-context example  $x_{i} \oplus y_{i} \oplus {}^{\prime \prime}\backslash n^{\prime \prime}$  is approximately independent from the other in-context examples:
$$
c _ {1} ^ {k} \leq \frac {\prod_ {i = 1} ^ {k} \mathbb {P} _ {\phi} \left(x _ {i} \oplus y _ {i} \oplus {} ^ {\prime \prime} \backslash n ^ {\prime \prime}\right)}{\mathbb {P} _ {\phi} (p)} \leq c _ {1} ^ {- k} \tag {8}
$$
In addition, Assumption 3 bounds the distribution drift that caused by the artificial new line token:
$$
\mathbb {P} _ {\phi} \left(x _ {i} \oplus y _ {i}\right) \leq \mathbb {P} _ {\phi} \left(x _ {i} \oplus y _ {i} \oplus {} ^ {\prime \prime} \backslash n ^ {\prime \prime}\right) \leq \mathbb {P} _ {\phi} \left(x _ {i} \oplus y _ {i}\right) \cdot c _ {2} \tag {9}
$$
Similarly, we can use Assumption 3 again to bounds also the distribution drift that caused by potential label flipping:
$$
c _ {2} <   \frac {\mathbb {P} _ {\phi} \left(x _ {i} \oplus \tilde {y} _ {i}\right)}{\mathbb {P} _ {\phi} \left(x _ {i} \oplus y _ {i}\right)} <   \frac {1}{c _ {2}} \tag {10}
$$
where  $\tilde{y}_i$  denote the "corrected" label. So overall, we got that:
$$
\left(c _ {1} c _ {2} ^ {2}\right) ^ {k} <   \frac {\prod_ {i = 1} ^ {k} \mathbb {P} _ {\phi} \left(x _ {i} \oplus \tilde {y} _ {i}\right)}{\mathbb {P} _ {\phi} (p)} <   \left(c _ {1} c _ {2} ^ {2}\right) ^ {- k} \tag {11}
$$
Now, we can write the ratio of the prompt probabilities according to the ground-truth component and the other mixture components explicitly and get that: