Understanding In-Context Learning via Supportive Pretraining Data
Xiaochuang Han, Daniel Simig, Todor Mihaylov, Yulia Tsvetkov, Asli Celikyilmaz, Tianlu Wang
Introduction
In-context learning in NLP has drawn tremendous attention recently (Dong et al., 2022). Unlike traditional learning paradigms that rely on training or finetuning models, in-context learning only provides a handful of demonstration examples to language models as a prefix to the test input, without any parameter updates. In-context learning has shown superior performance on a range of NLP tasks Brown et al. (2020); Zhang et al. (2022b); Chowdhery et al. (2022); Hoffmann et al. (2022), but the origin and reason of this emergent ability remain under-investigated. In-context learning is surprising since language models have not been explicitly trained to learn from demonstration examples (Xie et al., 2022). As shown in an illustrative scenario in Figure 1, a typical pretraining data instance is highly different from an in-context learning example for downstream tasks, in both content and format.
Prior work have attempted to answer what in-context learning is, through empirically investigating useful and irrelevant attributes of the demonstration examples (Min et al., 2022; Zhang et al., 2022a), or theoretically proving certain synthetic language models implicitly do Bayesian inference with demonstrations (Xie et al., 2022). Furthermore, recent work have drawn connections between the mechanism of in-context learning and standard learning algorithms, such as regression, nearest neighbor, and gradient descent (Olsson et al., 2022; Akyürek et al., 2022; Dai et al., 2022; von Oswald et al., 2022).
Differently, in this work we are interested in understanding from where the in-context learning ability is acquired, through a perspective of pretraining data. Although not many, some recent work have investigated this direction. For instance, Shin et al. (2022) pretrain a variety of language models on different corpora. They study correlations between attributes of pretraining datasets and in-context learning performance, at a relatively coarse dataset-level. Chan et al. (2022) construct pretraining data with different attributes and discover that some distributional properties of the data drive the emergence of in-context learning. However, their experiment is limited to synthetic data of image-label pairs.
In this work, we investigate a large language model OPT (Zhang et al., 2022b) and its pretraining data. We first hypothesize that there exists some specific pretraining data instances that are particularly helpful to the model’s in-context learning ability. As an attempt to find such instances, we adapt an iterative, gradient-based method ORCA (Han and Tsvetkov, 2022) to search within OPT’s pretraining corpus. The process is guided by the gradients of the in-context learning data from downstream tasks, and we refer to the identified subset as supportive pretraining data to in-context learning following Han and Tsvetkov (2022). Furthermore, we quantitatively verify through a perturbative continued pretraining, that the supportive subset does improve the model’s in-context learning performance on downstream tasks, while not affecting a spurious zero-shot performance (§2).
We then analyze the identified supportive data in contrast to the general pretraining data, to obtain data features particularly relevant to in-context learning. We specifically approach from three aspects: the domain relevance to downstream tasks, the token frequency distribution, and the information gain of incorporating long-range pretraining context. Our major findings include: (1) Compared to general pretraining data, the supportive data do not have a higher domain relevance to the downstream tasks. (2) The supportive pretraining data contain a relatively higher amount of rarely occurring, long-tail tokens. (3) The supportive pretraining data are challenging examples in incorporating long-range context for language modeling (§3).
Our work offers a first step towards interpreting in-context learning in NLP tasks via analyzing instance-level pretraining data. We believe it can help improve the transparency and interpretability of language models’ in-context learning behavior. Our analysis can also pave the way to improved in-context learning in the future by informing pre-training data construction.
Finding supportive pretraining data for in-context learning
Han and Tsvetkov (2022) propose an iterative, gradient-based method ORCA to find supportive pretraining data of BERT (Devlin et al., 2019) under a vanilla zero-shot prompting setup. In this section, we provide some background and adapt ORCA for large language models in a setting of in-context learning (ICL), finding supportive pretraining data for downstream tasks with demonstration examples.Identifying important training data for an inference time model output is an estabilished topic in model interpretability, with various prior work measuring data importance via variants of gradient similarity (Koh and Liang, 2017; Pruthi et al., 2020). However, these methods are prohibitively expensive to be applied to large-scale pretraining data. Concurrent to our work, Guu et al. (2023) propose an interesting method to model the importance of individual training examples by simulating training runs, but it is also on a scale of finetuning instead of pretraining.
Assume we have a pretrained language model (LM) and data pairs representing the inputs and ground truth outputs of task . Both and are in natural language. For classification tasks, the target labels can be converted to natural language via verbalizers (Schick and Schütze, 2021).
A pretrained language model can be applied to perform downstream tasks via zero-shot prompting (e.g., Petroni et al., 2019). For classification tasks, the language model outputs a candidate answer with top probability, , where contains all candidate answers . For generation tasks, outputs can be obtained by sampling autoregressively from conditioned on (e.g., Holtzman et al., 2019). This is a zero-shot scenario with no demonstration examples.
In-context learning
Instead of modeling , ICL estimates , prepending the original model input with several demonstration examples sampled from the target task . The language model is never trained on the task data with demonstrations. However, we can form a loss on the in-context data as a surrogate for ’s ICL performance, which will be used for a later guidance step, .
Pretraining
The pretraining data of often consists of texts from large, general-domain corpora. During pretraining, the LM is updated via stochastic gradient descent with a loss to reconstruct given a prefixing context, .
Supportive pretraining data
Our goal is to locate what pretraining data if upweighted would be most helpful to the LM ’s ICL ability. Following ORCA (Han and Tsvetkov, 2022), we use the similarity between gradients and iteratively to find such supportive pretraining data. We show details of our adapted algorithm ORCA-ICL in Figure 2. The algorithm finds pretraining data that exert a gradient to similarly as a group of guidance ICL task data would. provides a guidance for the direction the model parameters should be updated towards to be better at ICL, while approximates how the direction the model parameters would be updated based on individual pretraining instances. We conduct a multi-iteration process (a total of iterations each selecting supportive instances) to mitigate noise.Additionaly according to Han and Tsvetkov (2022), this may prevent selecting examples associated with only one class of the task, a case of poor calibration. SGD denotes an one-pass stochastic gradient descent to mimick an incremental upweight to the selected data, with a minimum number of steps to prevent overfitting. The resulting supportive set has a very small size (under 2000 in this work).More details of the ORCA algorithm can be found in Han and Tsvetkov (2022).
Verifying supportiveness
To quantitatively evaluate the supportiveness of the selected set of pretraining data, we perform an one-pass gradient descent on the original LM with the selected set , which mimics a perturbative continued pretraining with a minimum number of updates: . We then benchmark this perturbed model () with the original model () and a model perturbed with a random set of pretraining data. We expect the perturbed model using our selected supportive pretraining data to achieve a better ICL performance.
2 Setup
Throughout the work, we use a pretrained, autoregressive OPT-6.7B (Zhang et al., 2022b) as our LM .
Tasks
In this work, we focus on classification problems and first retrieve 48 classification-based tasks from Natural Instructions v2 (NI-v2, Wang et al., 2022). We apply the LM on the tasks with both a zero-shot and in-context learning setup. We extract tasks that achieve at least 10% better performance with in-context demonstrations. We group 17 tasks that satisfies the constraint and further select 6 typical tasks among them:
SST-2: Movie review sentiment classification (Socher et al., 2013). AG News: News topic classification (Zhang et al., 2015). Story Cloze Test: Story coherence classification (Mostafazadeh et al., 2017). SMS Spam Collection: Spam classification (Almeida et al., 2011). Sentiment 140: Tweet sentiment classification (Go et al., 2009). TweetQA: Answer verification (Xiong et al., 2019).
For each task, we randomly sample 500 examples with a balanced class distribution as , guiding the ORCA-ICL algorithm. The quantitative evaluation is performed on the full dataset. For ICL, for each instance in the task data, we randomly sample 4 demonstration examples under each candidate class defined in the task.The sampling of demonstration examples is independent across test instances to mitigate potential spurious correlations. The order of demonstration examples in the context is randomly shuffled. The template and verbalizer of each task follows the original NI-v2 dataset, though we did not include the task instructions, as the focus of this work is in-context learning with demonstration examples.
Pretraining
Considering the size of pretraining data , we include an as large portion of OPT’s pretraining data as possible under a reasonable budget. Specifically, in this work we use a total of 2.5M pretraining instances each consists of 2048 tokens.The total 5B tokens are about 3% of OPT’s 180B full pretraining data. For computing efficiency, we use intra-layer model parallelism (Shoeybi et al., 2019) and fully sharded data parallel (Ott et al., 2021).This groups 4 input data for each backward pass in our setup. The 4 instances receive a same gradient similarity score, equivalent to an aggregated instance 4 times of the length.
Implementation Details
We run ORCA-ICL with a maximum of iterations. In each iteration we extract pretraining instances with top gradient similarity with the ICL task data. We use a batch size of 16 and learning rate of 2e-5 for the one-pass gradient descent with an Adam optimizer (Kingma and Ba, 2014). This results in a total of 125 updatesThe one-pass descent has steps. to the original LM after all iterations as the perturbative continued pretraining.
3 Results
As the main evaluation of the supportive pretraining data obtained by ORCA-ICL, we perform perturbative continued pretraining on both the selected supportive data and random pretraining data as a control. Table 1 shows the main results of task accuracy. The leftmost column shows a source task guiding the selection of supportive pretraining data. At each row, we evaluate the perturbed model () on all 6 tasks. The ICL performance of the original LM is reported in the headers of the table.
In each cell of the table, the top number shows the continued pretraining result with the supportive data we identified. We consider iterations as a hyperparameter and report result with a best . We want to know at a same size of selection, how our identified subset performs compared to random pretraining data. We therefore run random selection with 5 seeds, and the bottom number of the cell shows the continued pretraining result with random data at a same size of our selection, accompanied by a standard deviation. The performance of our selection is bolded when the performance difference with random selection exceeds one standard deviation.
The diagonal cells show the performance of perturbed models on the same task used for selecting supportive data. We observe on 4 of the 6 source tasks, our selection of supportive pretraining data is effective. For the cross-task performance, we observe on 5 of the 6 source tasks, our selection is effective for at least three tasks.Negative result is observed with TweetQA, on which we conjecture the patterns in the demonstration examples are more difficult to transfer to the test input (e.g., factual knowledge instead of sentiment indicators). We conclude that our identified supportive pretraining data is overall effective for ICL, though the cross-task results show a portion of the ICL behavior can be task-specific and not universal across tasks.
Control evaluation on zero-shot data
Being effective on the ICL data does not necessarily mean a direct support for a model’s ICL ability, which is to learn from the demonstration examples. The test input can be a confounding factor: if our selection is effective as well on zero-shot test input without demonstrations, then the selection is not specific to the ICL ability. Therefore, we further confirm the supportiveness of our selected supportive pretraining data to ICL, contrastively in a zero-shot setup. We evaluate our models after perturbative continued pretraining in Table 1 on the same tasks but without the in-context demonstrations. We present the results in Table 2. The two columns show the zero-shot prompting performance of the original LM and the model after continued pretraining with our ICL-supportive selection, respectively. We do not observe performance gain for most tasks, indicating our selection is specific to the ICL ability without benefiting the zero-shot, no-demonstration task performance.
Analyzing supportive pretraining data for in-context learning
In the previous section, we identify a small subset of pretraining data that supports the ICL ability of language models. In this section, we analyze the selected supportive pretraining data to understand what makes them useful to ICL. Specifically, we compare the supportive pretraining data contrastively with randomly sampled pretraining instances, investigating three aspects of the pretraining data: the domain relevance to downstream tasks, the token frequency distribution, and the information gain of incorporating long-range context.
Xie et al. (2022) and Min et al. (2022) imply that in-context demonstration is useful since it helps locate a particular domain or concept of the test input the LM already learned through the pretraining data. On the other hand, Olsson et al. (2022) imply that in-context demonstration is useful because the decision over the test input may be done through a soft-copy mechanism from the demonstration examples. These lead to two different expectations of the role of supportive pretraining data: (1) Inferred from Xie et al. (2022) and Min et al. (2022), the supportive pretraining data should be from a same domain as the demonstration and test examples, providing direct supporting knowledge to solve the downstream task. (2) Inferred from Olsson et al. (2022), the supportive pretraining data should be beneficial to the soft-copy mechanism, providing meta support for the abstract ability, unconstrained with the concrete data domain.This view of supportive data will be revisited in §3.3. We aim to measure the domain relevance between supportive pretraining data and downstream tasks.
To quantify domain relevance, we use MAUVE score (Pillutla et al., 2021) to measure an information divergence between two text distributions. We compute two MAUVE scores, between the target task data and our selected supportive pretraining data, and between the task data and random pretraining data. We then compute and report their difference. A positive MAUVE difference indicates a higher domain relevance of our supportive pretraining data.Pillutla et al. (2021) also shows higher MAUVE indicates higher generation quality, but we skip that aspect since all of our data are naturally occuring text. We use RoBERTa (Liu et al., 2019) as MAUVE’s embedding model following He et al. (2022).
Results
We show the difference of MAUVE scores in Figure 3. The error bar shows the 95% confidence interval using 32 random seeds. We find that for 5 of the 6 tasks, there is no significant difference between the MAUVE scores of supportive pretraining data and random data. For SST-2, the supportive pretraining data even shows a lower MAUVE score. Therefore, the supportive pretraining data to ICL do not have a higher domain relevance to the task, compared to general pretraining data. This result aligns with the domain relevance finding in Shin et al. (2022) where dataset-level analyses were performed. This implies the improved ICL behavior of our models may be a meta ability, aided by pretraining data unrelated to the specific domain knowledge for solving the task, but related to a domain-invariant mechanism to learn from a data’s context. §3.3 continues this discussion.
2 Token frequency distribution
Providing demonstrations to a task input under an ICL setup creates repetitions (e.g., of label tokens), which changes the token frequency distribution of the ICL task data. Therefore, we are interested in whether the supportive pretraining data possess a different token frequency distribution from general pretraining data. Experimented with sequences of image-label pairs, Chan et al. (2022) find that a skewed class distribution (high burstiness) and a large number of rarely occurring classes in training data promote the ICL ability of Transformer models (Vaswani et al., 2017). However, it is unknown whether the findings on the synthetic image-label data can transfer to the natural language pretraining data, a gap we address in this subsection.
We fit a Zipfian distribution over each supportive and random pretraining instance that consists of 2048 tokens. The Zipf’s coefficient is the negative slope of a linear regression over the tokens’ log-rank v.s. log-frequency. A higher Zipf’s coeffcient indicates a higher mass on the frequent tokens (i.e., more skewed distribution). A lower Zipf’s coefficient indicates a higher mass on the rare, long-tail tokens (i.e., flatter distribution).
Results
In Figure 4, we show the difference in average Zipf’s coefficients between supportive and random pretraining data, each with a group size of 2000. The error bar shows the 95% confidence interval with 32 random seeds. We find that for all tasks, the Zipf’s coefficient of the supportive pretraining data is significantly lower than that of the random pretraining data. This indicates a flatter Zipfian distribution with a relatively higher mass over the long-tail tokens. In other words, though the overall burstiness of data is lower, there is a relatively higher amount of rarely occurring, long-tail tokens in the supportive pretraining data for ICL. Flatter frequency distribution also indicates higher entropy over the tokens, presumably making the supportive pretraining data challenging examples to fit by the model, a concept we explore further in the next subsection.
3 Information gain from long-range context
In §3.1, we find that the domain relevance of the supportive pretraining data to downstream tasks is not higher than that of random pretraining data. This is comprehendible if we follow the aforementioned perspective of Olsson et al. (2022), hypothesizing that there exists a soft-copy mechanism between the in-context demonstrations and test input. The supportive pretraining data may provide meta support for the abstract soft-copy mechanism rather than task-specific knowledge. We further hypothesize that to facilitate such meta support, the incorporation of long-range context during language modeling in supportive pretraining data should be different from random pretraining data, since the demonstration examples in the ICL setup is a form of long-range context. We propose a novel information gain measure to quantify this feature of incorporating long-range context.
Ideally the length of long or short context should remain constant across different tokens , but it would be a very expensive computation due to a lack of parallelism. We approximate the computation by splitting a full sequence of pretraining tokens (e.g., 2048 tokens) to smaller blocks and calculate cross entropy with the boundary of blocks:
With the above definition, the average length of context for all is and , respectively. In the experiments below, we keep for the length of short context and increase the length of long context at .
We report the difference in the average information gain (across ) of incorporating long-range context for a language modeling objective, in supportive pretraining data over random pretraining data. Additionally, we want to use the defined information gain measure as a standalone feature of data, so we use a different LM to compute the cross entropy than the LM on which we perform ICL. Below we report results using OPT-1.3B, while experiments using OPT-350M shows a similar trend.
Results
In Figure 5, we see for all of the experimented tasks, there is a significant trend that increasing the length for the long-range context for supportive pretraining data has a lower relative information gain compared to random pretraining data. Though seeming counterintuitive at first glance, this suggests that the supportive pretraining data are more challenging examples in incorporating the long-range context information.Note that a reverse of the statement may not hold necessarily, since an example’s long-range context can also be irrelevant by nature and challenging in a useless way. A possible explanation for this is that such challenging examples contain confounding spans that harms the information gain measure. The language model has to learn to decide which part of the long-range context is truly relevant to the prediction of next tokens. This would resemble more and thus helpful to the ICL task scenario where there are multiple demonstrations from different classes.
4 Future work
Despite our aforementioned findings, we mainly conduct correlational analyses throughout the work. Despite the potential confounding factors, future work can try converting the correlational findings to causal ones. For example, to actively refine or construct pretraining data to improve existing models’ ICL performance, with a metric of token frequency distribution (i.e., find data with a higher mass of long-tail tokens) or context information gain (i.e., find difficult examples in incorporating long-range context). Additionally, we only investigate classification tasks in this work. However, the ORCA-ICL method can be applicable to generation tasks as well in the future, if the ICL loss is defined over a sequence probability of the generation.
Related Work
Min et al. (2022) understand ICL through analyzing which aspects of the demonstration examples contribute or are irrelevant to task performance. They find replacing ground truth demonstration labels with random labels would not hurt task performance, while ICL still benefits from knowing the label space, distribution of inputs, and sequence format specified in demonstration examples.Recent work like Wei et al. (2023) and Pan et al. (2023) show the related findings would depend on model scales as well. Zhang et al. (2022a) further show on sequence labeling tasks, the length of demonstrations and the relevance of their tokens are important for ICL.
Learning mechanism
Xie et al. (2022) explain ICL as implicit Bayesian inference, occurring when language models infer a shared latent concept from demonstration examples at inference time. They show language models exhibit such ICL behavior by constructing synthetic pretraining data with a controlled distribution of concepts. Garg et al. (2022) empirically show that Transformer models can be trained to learn unseen linear functions from in-context demonstration examples. Olsson et al. (2022) present evidence that multi-layer attention-based models form an induction head and perform ICL by a pattern copying behavior from the prefixing context. More recent work like Akyürek et al. (2022), Dai et al. (2022), and von Oswald et al. (2022) explain ICL in Transformer models as a kind of standard learning algorithms over the demonstration examples, such as gradient descent and regression.
Pretraining data
Razeghi et al. (2022) find on numerical reasoning tasks, a language model’s ICL performance is highly correlated with the term frequency of the input data in the pretraining corpus. Shin et al. (2022) investigate how ICL can be affected when the pretraining dataset varies. They discover that ICL heavily depends on the corpus domain source, but pretraining with a corpus related to a downstream task does not always translate to a competitive ICL performance on the task. Chan et al. (2022) experiment on a synthetic image-label pairs dataset. They show certain distributional properties of the synthetic pretraining data, such as the burstiness of classes and large numbers of rarely occurring classes, promote the emergence of ICL. Our work belongs to this line of work, but offers a first step towards understanding ICL in realistic NLP tasks through analyzing instance-level pretraining data. Additionally, concurrent to our work, Gu et al. (2023) propose a method that groups pretraining data by their instrinsic tasks, enhancing instead of interpreting existing language models’ ICL ability.
Conclusion
In-context learning has shown superior performance on a range of NLP tasks, yet it remained unclear from where language models acquired this ability. We approach the problem by identifying a small subset of pretraining data that particularly supports language models to do in-context learning on downstream tasks. We analyze common features of the supportive instances in contrast to general pretraining data and find that: (1) The supportive pretraining data do not have a higher domain relevance to the downstream tasks. (2) The supportive data contain a relatively larger amount of rare, long-tail tokens. (3) The supportive pretraining data are more challenging instances in incorporating long-range context in language modeling. Our findings may be beneficial to future work that refine or construct pretraining data, in order to actively improve existing models’ in-context learning performance.
Limitations
It is worth noting that the supportive pretraining data we investigated throughout the work is w.r.t. the current LM, such that a perturbative continued pretraining with the supportive data would improve the final LM checkpoint deployed to downstream tasks. It is possible that for some data which we did not determine as supportive, they had been supportive w.r.t. early checkpoints of the LM. With more computing resources, future work may investigate the trend of supportive patterns across multiple checkpoints of a LM throughout the pretraining process.
Additionally, another significant limitation of our work is the amount of involved computing resource. The ORCA-ICL method is gradient-based that requires back-propagation. Since we iterate through a large size of pretraining data, the cost of computation is similar to training a language model with a batch size of 1 on the considered pretraining data. On our 4 nodes each consists of 8 Nvidia V100 GPUs, finding the supportive pretraining data for each source task in our experiment would take about a week. One mitigating aspect of such computation is that the gradient calculation can be done asynchronously, therefore enabling the use of idle, leftover GPUs scattered across a cluster of nodes. We plan to explore efficient computation of gradient similarity or move from a paradigm of extracting supportive data to generating supportive data in future work.
Acknowledgements
We thank Naman Goyal, Anjali Sridhar, Zeyu Liu, Victoria Lin, Mengzhou Xia, Weijia Shi, Jiacheng Liu, Hao Zhu, and Tianxing He for helpful discussions. We also thank the anonymous ACL reviewers and all members of TsvetShop for the valuable feedback. This research is supported in part by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via the HIATUS Program contract #2022-22072200004. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of ODNI, IARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.
References
Appendix A Qualitative examples
In Appendix A, we show some qualitative examples of the supportive pretraining data to ICL and random pretraining data. Note that these are illustrative examples extracted from long pretraining instances (each instance consists of 2048 tokens), for a better understandability of our findings. A manual examination of such data is difficult, and we thus propose the quantitative analyses described in the main paper.