Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale (2023)

Abstract: Language models have been shown to perform better with an increase in scale on a wide variety of tasks via the in-context learning paradigm. In this paper, we investigate the hypothesis that the ability of a large language model to in-context learn-perform a task is not uniformly spread across all of its underlying components. Using a 66 billion parameter language model (OPT-66B) across a diverse set of 14 downstream tasks, we find this is indeed the case: ~70% of the attention heads and ~20% of the feed forward networks can be removed with minimal decline in task performance. We find substantial overlap in the set of attention heads (un)important for in-context learning across tasks and number of in-context examples. We also address our hypothesis through a task-agnostic lens, finding that a small set of attention heads in OPT-66B score highly on their ability to perform primitive induction operations associated with in-context learning, namely, prefix matching and copying. These induction heads overlap with task-specific important heads, reinforcing arguments by Olsson et al. (2022) regarding induction head generality to more sophisticated behaviors associated with in-context learning. Overall, our study provides several insights that indicate large language models may be under-trained for in-context learning and opens up questions on how to pre-train language models to more effectively perform in-context learning.

## introduction in recent years, large language models (llms) @xcite @xcite @xcite @xcite based on the transformer architecture @xcite pre-trained using self-supervision on web-scale textual corpora have revolutionized the field of natural language processing (nlp). at larger scales, these models demonstrate remarkable emergent @xcite prowess in performing a wide variety of tasks without any form of fine-tuning, via the zero/few-shot incontext learning paradigm @xcite . how in-context learning works has been an open question since its advent and recent studies @xcite @xcite have begun scratching the surface toward better understanding the paradigm. in this paper, we empirically address the following key question: are all llm components really needed to perform in-context learning? the first way we address the aforementioned question is through the lens of task-specific importance scores and structured pruning @xcite @xcite of components underlying modern llms, which are primarily stacks composed of multiple highdimensional self-attention blocks that form multiheaded attention and densely activated feed forward networks (ffns). we pick the open pretrained transformer (opt) @xcite model with 66b parameters for our analyses, which yield several surprising observations. we find that important attention heads are primarily clustered in the intermediate layers and important ffns are primarily in later layers of the model ( §4). we find that the ability to perform zero/few-shot in-context learning on a variety of 14 nlp datasets/tasks stays nearly intact when up to 70% (∼15.7b parameters in opt-66b) of the attention heads are removed ( §5.1). the attention heads that are (un)important for in-context learning also seem to overlap across tasks ( §6.1) and shots ( §6.2), and pruning attention heads based on a "universal" importance order computed using all 14 datasets generalizes to varying degrees on out-of-distribution datasets ( §6.1.2). these observations indicate that a common taskagnostic subset of the attention heads are responsible for in-context learning. we also find that only up to 20% of the ffns (∼8.5b parameters) can be removed with minimal decline in zero/few-shot in-context learning performance ( §5.2), indicating the importance of ffns toward in-context learning. the second way we address the aforementioned question is by quantifying the capacity of all attention heads in opt-66b to perform a subset of task-agnostic primitive operations associated with in-context learning, namely, prefix matching and copying: explicitly searching for a prior occurrence of the current token in-context and copying over its suffix. @xcite and @xcite developed a mathematical framework to reverse-engineer a transformer and also find such heads, termed induction heads, and explored the hypothesis that such heads drive in-context learning with model sizes up to 13b parameters in a mostly task-agnostic fashion. using this framework, we compute task-agnostic scores for prefix matching and copying for each attention head and find that a small set of heads in opt-66b have nontrivial scores for both primitives ( §6.3). qualitative inspection and quantitative analyses show that these heads overlap (to varying degrees) with the ones identified earlier to be important for in-context learning via our set of 14 nlp datasets/tasks, indicating that induction heads are capable of more sophisticated behaviors associated with in-context learning such as latent concept matching but are not the only heads with such capabilities ( §6.3.1). overall, our study provides several insights about in-context learning at massive scale using both task-specific and task-agnostic settings. in a world of ever increasing language model sizes, we believe these insights serve as a strong foundation for researchers and practitioners in language modeling to build and leverage compact language models that can also demonstrate emergent abilities. ## background & methods in this section, we establish notation and methods with the open pre-trained transformer (opt) @xcite model used for our study, provide background on in-context learning and the mathematical formulation of induction heads by @xcite that we build on, and describe our adaptation of oracle and gradient-based importance score formulations for in-context learning. ## experimental setup we conducted our experiments on opt-66b, which was the largest publicly available dense decoderonly language model at the time of our experiments. we efficiently compute gradient-based importance scores for the 4608 attention heads and oracle importance scores for the 64 feed forward networks (ffns) in opt-66b. we experiment with a variety of 14 nlp datasets/tasks. for consistency in the evaluation metric, we report accuracy on all tasks. our choice of datasets and metric is in line with @xcite . the datasets include arc easy and challenge @xcite and openbookqa @xcite for advanced question-answering, hel-laswag @xcite , piqa (bisk et al., 2020) and winogrande @xcite for various forms of commonsense reasoning, and the following datasets from the standard superglue benchmark ( wang et al., 2019 figure 3 depicts a heatmap of the head importance scores averaged across all tasks (as described in §2.3.2) in the 5-shot setting. task-averaged heatmaps for the 0-shot and 1-shot settings and all task-specific heatmaps are provided in appendix a.1. we observe that the important attention heads are primarily clustered in the intermediate layers of opt-66b in both the task-averaged and taskspecific cases. we also observe overlap in the important heads across the different zero/few-shot settings, confirmed in follow-up analysis in §6.2. ## iterative pruning we now assess to what extent we can remove multiple attention heads and/or ffns with minimal decline in task performance. for each task in each (0/1/5-shot) in-context learning setting, we create separate rankings of attention heads and ffns in opt-66b by separately sorting them in ascending order by their importance scores ( §4.1 and §4.2). we then remove unimportant attention heads or ffns in an iterative fashion using these rankings, 10% at a time, and re-evaluate task performance after each removal. 1 ## detailed analysis of attention heads in this section, we perform a detailed analysis of the attention heads in opt-66b, given that in-context learning is auto-regressive in nature and attention heads explicitly encode cross-token interactions. ## related work there has been an interest in effectively leveraging the in-context learning paradigm @xcite @xcite @xcite @xcite ever since its introduction by brown et al. (2020), but there have been relatively fewer studies toward better understanding the paradigm itself. xie et al. (2021) cast in-context learning as implicit bayesian inference where the language model implicitly infers a shared concept among in-context examples when making a prediction. min et al. (2022b) study the role of the in-context examples themselves, finding that the ground-truth labels are not needed in the examples and that the more important drivers are provision of the label space, the distribution of the input text and the overall format of the sequence. garg et al. @xcite showcase that transformer models trained from scratch can in-context learn the class of linear functions with performance comparable to the optimal least squares estimator even under distribution shifts. @xcite showcase that incontext learning performance is correlated strongly with term frequencies in the pre-training corpora used. @xcite consider an alternate framing of in-context learning as the ability of a language model to better predict tokens later in the context than tokens earlier and hypothesize the existence of induction heads that are responsible for in-context learning. @xcite show that transformers exhibit striking differences in generalizing from in-context vs. in-weights information. several works have also focused on analyzing and interpreting how attention works. @xcite performed a study on gpt-2, finding that attention targets different parts of speech at different layer depths and aligns with dependency relations most strongly in the middle layers. @xcite showcase that bert encodes the classical nlp pipeline in an interpretable way across layers. there are works relying on different formulations for head importance, such as layerwise relevance propagation @xcite , gradient-based importance and oracle knock-off importance @xcite , with small taskspecific trained models and report the existence of specialized heads. given the recent trend of increasing model scale @xcite @xcite toward tuning-free general-purpose language models that exhibit emergent in-context learning abilities, we draw and build on prior work to understand just how much scale is really needed and/or used for in-context learning downstream, an aspect somewhat eclipsed by the focus on the pre-training loss curve in scaling laws @xcite . it is also worth noting that some of our empirical observations rely on a simple greedy approach to training-free pruning since our focus was not to optimally prune a language model with respect to performing in-context learning. @xcite show the greedy approach is sub-optimal and produces under-estimates and halabi et al. ( 2022 ) account for the need to re-compute importance scores after removal of each attention head or ffn by formulating pruning as weakly sub-modular maximization. ## conclusion & future work in this paper, we studied the efficacy of attention heads and feed forward networks (ffns) in a large language model (opt-66b) in performing in-context learning in both task-specific and task-agnostic settings. we observed that while in-context learning may have emerged via selfsupervised pre-training at scale, only a core nucleus of attention heads and ffns seem to be important for in-context learning across a wide variety of downstream tasks. we observed that a small set of attention heads have the capacity to perform task-agnostic primitive induction operations associated with in-context learning, namely, prefix matching and copying. we also saw that these induction heads overlap with task-specific important attention heads, indicating that induction heads are capable of more sophisticated forms of incontext learning and reinforcing arguments @xcite about their generality. overall, our incontext learning-centric observations complement recent work @xcite in indicating that large language models may be under-trained and motivate several interesting directions for future work. while induction heads are formed naturally during self-supervised pre-training in its current form, we believe it may be possible to increase the number and strength of induction heads formed by defining auxiliary pre-training objectives for primitives like prefix matching and copying. more generally, it may also be prudent to investigate and improve (pre-)training regimes to increase the number of important model components to in-context learn-perform a wide variety of downstream tasks. multi-task instruction-tuning likely belongs to this category and it would be interesting to replicate our study with now increasingly accessible instructiontuned model variants (such as opt's instruction meta-learned variant opt-iml).