Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale
Year: 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 incontext 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 incontext 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 undertrained for in-context learning and opens up questions on how to pre-train language models to more effectively perform in-context learning.

heading: 1 Introduction
text: In recent years, large language models (LLMs) (Brown et al., 2020; Rae et al., 2021; Lieber et al., 2021; Black et al., 2022; Zhang et al., 2022; Chowdhery et al., 2022; Hoffmann et al., 2022; Smith et al., 2022) based on the Transformer architecture (Vaswani et al., 2017) pre-trained using self-
∗Work done as an intern at AWS AI Labs. †Corresponding author.
Code: github.com/amazon-science/llm-interpret
supervision on web-scale textual corpora have revolutionized the field of natural language processing (NLP). At larger scales, these models demonstrate remarkable emergent (Wei et al., 2022) prowess in performing a wide variety of tasks without any form of fine-tuning, via the zero/few-shot incontext learning paradigm (Brown et al., 2020).
How in-context learning works has been an open question since its advent and recent studies (Xie et al., 2021; Garg et al., 2022; Olsson et al., 2022; Min et al., 2022b) 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 (Li et al., 2016; Molchanov et al., 2016; Anwar et al., 2017) 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) (Zhang et al., 2022) 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
11833
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. Elhage et al. (2021) and Olsson et al. (2022) 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.

heading: 2 Background & Methods
text: In this section, we establish notation and methods with the Open Pre-trained Transformer (OPT) (Zhang et al., 2022) model used for our study, provide background on in-context learning and the mathematical formulation of induction heads by Olsson et al. (2022) that we build on, and describe our adaptation of oracle and gradient-based importance score formulations for in-context learning.

heading: 2.1 Open Pre-trained Transformer (OPT)
text: OPT is a suite of language models of varying sizes aimed at serving as open replicas of GPT-3. The largest openly accessible model from this suite is OPT-66B with 66 billion parameters.
Architecture: Consider a tokenized input sentence to OPT, X ∈ RN×de , where N is the number of tokens in the sentence and de is the embedding dimension. The input is processed by multiple decoder layers consisting of multi-headed attention (MHA) blocks, layer norm (LN) and feed forward networks (FFN), followed by a linear layer to produce logits over the vocabulary. The decoder layers can be formally expressed as follows:
t(ℓ+1) = zℓ + MHAℓ(LNℓ(zℓ)) (1)
z(ℓ+1) = t(ℓ+1) + FFNℓ(t(ℓ+1)) (2)
where z1 = X, and (1) & (2) are the residual connections corresponding to the MHA and FFN in layer ℓ >= 1 respectively. OPT-66B was pre-trained with a maximum sequence length of 2048 and embedding dimension de = 9216.
MHA: In an MHA block, H attention heads are applied in parallel to the input and their outputs are concatenated. In OPT-66B, there are H = 72 attention heads of dimension dh = 128 in every layer ℓ. An individual attention head h in layer ℓ consists of three learnable matrices, Whk ,W h q ,W h v ∈ Rde×dh , all unique to the head, such that it applies selfattention Ah(.) on the input, where dh = de/H . Formally, for input M in layer ℓ:
MHAℓ(M) = [A1(M); · · · ;AH(M)]Wℓo (3) Ah(M) = sh(M)MWhv (4)
sh(M) = σ (MWhq (Whk)TMT√
dh
) (5)
where σ is the softmax function and Wℓo ∈ Rde×de is a learnable output matrix unique to the MHA
block in layer ℓ. To ensure OPT is auto-regressive, the output of sh(.) is masked to prevent the dependence of the hidden state of the token i, zℓi ∈ Rde , on future tokens in indices {i+ 1, . . . , N}.
To remove a head h in layer ℓ in practice, we set Ah(M) to be the zero matrix in Equation (3). This implies that Whk ,W h q ,W h v can be entirely removed, and the corresponding dh rows in Wℓo can also be removed. In total, there are 4608 attention heads across 64 layers in OPT-66B that constitute 21.7B of the total 66B parameters.
FFN: Each layer ℓ consists of a feed forward network (FFN) parameterized by a high-dimensional projection matrix, Wℓ1 ∈ Rde×d followed by a low-dimensional projection matrix, Wℓ2 ∈ Rd×de where d = 36864 for OPT-66B. Formally, for input M in layer ℓ:
FFNℓ(M) = ReLU(LNℓ(M)Wℓ1)W ℓ 2 (6)
where ReLU is the rectified linear unit activation function and LN is the layer norm.
To remove an FFN in layer ℓ in practice, we set FFNℓ(M) to be the zero matrix in Equation (6). This implies Wℓ1, W ℓ 2 and the layer norm LN
ℓ(.) for the FFN can be entirely removed. In total, FFNs constitute 43.4B parameters in OPT-66B.

heading: 2.2 In-Context Learning & Induction Heads
text: With increasingly larger language models being trained in recent years, a new paradigm of learning termed in-context learning (Brown et al., 2020) has become popular. In this paradigm, language models perform tasks by being prompted to generate output text conditioned on a few (or zero) incontext training examples that form solved "inputoutput" pairs for the task along with a query input. Figure 1 illustrates the paradigm for the task of identifying the sound that an animal makes. In some cases, tasks can also be accompanied by task descriptions/templates to help prime the language model better, e.g., zero-shot translating from English to German using the prompt:
English phrase: I like dogs. German phrase:
While these examples involve learning and relying on latent concepts during inference, few-shot in-context learning can additionally involve explicit primitive interactions between the in-context examples. For example, with the prompt:
English phrase: I like dogs. German phrase: ich
mag Hunde.
English phrase: I like ducks. German phrase:
the model may rely on prior in-context translations of the tokens I and like when performing the task for the query input. Olsson et al. (2022) developed a mathematical framework toward better understanding such mechanics, starting off with a task-agnostic formulation of in-context learning as the ability of a model to better predict tokens later in the context than the tokens earlier (Kaplan et al., 2020). They define a set of task-agnostic primitive operations that reflect the kind of interactions we refer to in the above example, namely, prefix matching and copying. These operations are defined in a simplistic fashion on a repeated sequence of randomly generated tokens: explicitly searching for a prior occurrence of the current token in-context and copying over its suffix. The heads that are capable of performing these operations are termed induction heads. Figure 2 depicts these operations for a repeated sequence of tokens. While these operations are intertwined in practice, the capacity of attention heads to independently perform them is computed with the scoring algorithms described in detail in Appendix A.8.

heading: 2.3 Importance Scores
text: Consider a modelM and a dataset D = {X ,Y}, where X = {x1, · · · ,xL} and Y = {y1, · · · ,yL} such that xi represents a prompt with few (or zero) in-context training examples along with a query input and yi represents the corresponding target output sequence. We define and compute importance scores for model components using such datasets to quantify their relative contributions to the model’s ability to perform in-context learning.

heading: 2.3.1 Oracle
text: Let PM(D) denote a dataset/task-specific performance metric, e.g., accuracy. Given dataset D, the oracle importance score of a component C inM is
computed as follows:
ISC(D) = PM(D)− PM\C(D) (7)
whereM\C denotes the resultant model when C is pruned fromM. Clearly, if pruning a component leads to poor model performance on the task, it must be important for the task. Similarly, if there is no difference or an improvement in performance upon pruning a component, it must be unimportant. Computing oracle importance scores for K model components requires us to perform O(K) evaluations for each dataset D.

heading: 2.3.2 Gradient-based
text: Given dataset D, the gradient-based importance score (Molchanov et al., 2016; Michel et al., 2019) of an attention head h captures the expected sensitivity of the model to h and is computed as follows:
ISh(D) = E(x,y) ∣∣∣∣Ah([x;y])T ∂L(y|x) ∂Ah([x;y]) ∣∣∣∣ (8)
where ; is the concatenation operator, (x,y) ∼ D such that x is a sequence of Tx tokens x1:Tx , y is a sequence of Ty tokens y1:Ty , A
h is the output of head h defined in (4) and the loss term in (8) is computed using the auto-regressive decomposition of the log-likelihood:
L(y|x) = − 1 Ty
j=Ty∑
j=1
log(p(yj |x, y1:j−1)) (9)
These importance scores can be efficiently computed for all heads by simply performing a single forward and backward pass over the model with D.
We also define the aggregated importance score of an attention head on a set of datasets S = {D1, · · · ,DK} as follows:
ISh(S) = ED∼S [ISh(D)] (10)

heading: 3 Experimental Setup
text: 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 Zhang et al. (2022). The datasets include ARC Easy and Challenge (Clark et al., 2018) and OpenBookQA (Mihaylov et al., 2018) for advanced question-answering, HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020) and Winogrande (Sakaguchi et al., 2021) for various forms of commonsense reasoning, and the following datasets from the standard SuperGLUE benchmark (Wang et al., 2019): BoolQ, CB, COPA, MultiRC, ReCoRD, RTE, WiC, and WSC. For a subset of experiments involving evaluation of outof-distribution generalization, we also use 2 additional datasets: MathQA (Amini et al., 2019) and LAMBADA (Paperno et al., 2016). We use a modified version of the lm-evaluation-harness framework (Gao et al., 2021) for our experiments. The default framework samples in-context examples at random, which we use without modification.

heading: 4 Importance Scores for OPT-66B

heading: 4.1 Attention Heads
text: 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.

heading: 4.2 Feed Forward Networks
text: We compute oracle importance scores (both taskspecific and averaged across tasks) for each FFN as described in §2.3.1 in the zero/few-shot settings.
Figure 4 depicts the scores in the 5-shot setting. The 0/1-shot scores are depicted in Appendix A.2.
We observe in the 0/1-shot settings that the removal of any FFN in the early (1-30) layers of OPT-66B either gives comparable or better performance for a vast majority of tasks. In the 5-shot setting however, both the early and later layers seem to have important FFNs for most tasks. We also generally observe high variance in FFN importance scores in later layers. We particularly note high variance for WSC and MultiRC, observing that removal of some individual FFNs can lead to absolute accuracy improvements/degradation of up to 20%! We leave further investigation into the cause for this variance for future work.

heading: 5 Iterative Pruning
text: 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

heading: 5.1 Removing Attention Heads
text: Figure 5 depicts the resulting task-specific and task-averaged accuracy trends in the 5-shot setting. Corresponding 0/1-shot trends are depicted in Appendix A.3. We observe that the average accuracy across tasks does not change much up until
1We do not remove attention heads one at a time and reevaluate given the number of heads and evaluation cost.
∼70% of the attention heads are removed. A finegrained look at the individual tasks also mostly shows similar trends, with accuracy staying fairly intact until a large proportion of the heads are removed. Some oddities include tasks such as WSC and CB, wherein we see that the 0-shot accuracy actually increases after removal of 70% of the heads.

heading: 5.2 Removing FFNs
text: Figure 6 depicts the resulting task-specific and task-averaged accuracy trends in the 0-shot setting. Corresponding 1/5-shot trends are depicted in Appendix A.4. We observe that in the 0-shot setting, the average accuracy across tasks does not change up until ∼20% of the FFNs are removed. For some tasks such as PIQA, Winogrande and RTE, the accuracy does not change even if 30% of the FFNs (∼13B of the 66B parameters) are removed. We also observe that the inflection point after which we observe a sharp decline in accuracy changes to 10% for the few-shot settings. Overall, these observations indicate that FFNs play a critical role toward in-context learning.

heading: 5.3 Combined Removal of Heads & FFNs
text: We now investigate whether the inflection points to in-context learning performance when removing either attention heads or FFNs in an iterative fashion still hold when removing them in tandem. Figure 7 depicts the average 5-shot accuracy of all tasks on joint iterative removal of attention heads and FFNs. Corresponding 0/1-shot trends are depicted in Appendix A.5. We observe that the removal of 70% of the attention heads (∼15.7B parameters) and 20% of the FFNs (∼8.5B parameters) leads to a mere 5% absolute drop in the average 0-shot accuracy. In the 1-shot setting, the drop in accuracy is 6% on removing 70% of the attention heads and 10% of the FFNs. In the 5-shot setting, the drop in accuracy is 4% on removing 60% of the attention heads and 20% of the FFNs. Overall, these new inflection points have deviated by at most 10% absolute, which may be attributed to the interplay between heads and FFNs.

heading: 6 Detailed Analysis of Attention Heads
text: 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.

heading: 6.1 Cross-Task Analysis
text: Michel et al. (2019) found preliminary empirical evidence of the existence of "universally" important attention heads in trained task-specific Transformer and BERT models via evaluating on out-of-domain test sets for machine translation and natural language inference respectively. With similar motivation, we study if the (un)important attention heads identified in various in-context learning settings for OPT-66B are shared across tasks.

heading: 6.1.1 Spearman’s Rank Correlation
text: We assess overlap in (un)important attention heads across tasks by sorting task-specific head importance scores to get head importance rankings and computing the Spearman’s rank correlation coefficient (SRCC) between the rankings for every pair of tasks in the zero-shot and few-shot settings. We also sort the task-aggregate head importance scores to get the aggregate ranking and compute the SRCC against the ranking for every constituent task. All correlations are depicted in Figure 8 for the 5-shot setting and Appendix A.6 for the 0/1-shot settings.
In both zero and few-shot settings, we observe statistically significant (p < 0.01) positive correlations in the head importance rankings for every pair of tasks, as well as between every task’s ranking and the aggregate ranking. This indicates that the set of (un)important attention heads are clustered together across tasks. We also observe seemingly lower magnitude SRCC values between every task and ReCoRD, a long reading comprehension task which requires commonsense reasoning, indicating the amount of head overlap is proportionally lower.

heading: 6.1.2 Generalization Trends
text: To understand how well head importance rankings generalize across tasks, we study accuracy trends for tasks when pruning using various head importance rankings. We study two sets of tasks.
The first set of tasks we study were used to compute the aggregate ranking: COPA, Winogrande and ReCoRD. For each of these 3 tasks, we consider the impact of pruning based on the self-ranking, aggregate ranking and the rankings from the tasks which share the highest and lowest SRCC with them. Figures 9a, 9b and 9c de-
pict the accuracy trends for these 3 tasks in the 5-shot setting. Corresponding trends in the 0/1- shot settings are in Appendix A.7. In the 0-shot setting, we observe that the accuracy on all 3 tasks when pruning using the rankings described is almost unaffected up to the 50% mark. We then observe a sharp decline in accuracy on COPA and Winogrande when the model is pruned to the 70% mark using the ranking identified via ReCoRD, the task with the lowest SRCC (0.13) with both COPA and Winogrande. This indicates that even if the rankings vary between ReCoRD and COPA/Winogrande (as reflected in the low magnitude of the SRCC score), the set of attention heads important for 0-shot learning with ReCoRD are important for COPA/Winogrande too. To further verify this, we calculated and found 71% and 76% overlap between the top 30% important attention heads for ReCoRD-COPA and ReCoRDWinogrande respectively. Comparing the zero-shot setting against the few-shot settings, we note that the decline/divergence in accuracy beyond the 50% pruning mark using the ReCoRD ranking is less sharp for COPA and Winogrande in the 1-shot setting and fades away in the 5-shot setting, indicating a convergence of important heads across tasks.
The second set of tasks we study are unseen, i.e., not used to compute the aggregate ranking: MathQA and LAMBADA. For these tasks, we analyze accuracy trends when pruning using the selfranking and aggregate ranking. Figures 9d and 9e depict their accuracy trends in the 5-shot setting. Corresponding trends in the 0/1-shot settings are in Appendix A.7. As expected, we observe that the self-ranking accuracy curves are somewhat higher than the aggregate ranking accuracy curves in general across both tasks. For MathQA, we also observe that the absolute difference in accuracy for both cases is within 1-2%. These indicate that the aggregate rankings generalize well to MathQA but not as much to LAMBADA.

heading: 6.2 Cross-Shot Analysis
text: To see if the attention heads identified to be (un)important for a task are shared across the different zero and few-shot settings, we compute Spearman’s rank correlation coefficient (SRCC) between the cross-shot head importance rankings for each task and compute the mean and variance across all 14 tasks. We observe that the mean SRCC is higher for rankings within the few-shot setting (0.41 for
1-shot vs. 5-shot) than for rankings across the zero and few-shot settings (0.39 for 0-shot vs. 1-shot and 0.37 for 0-shot vs. 5-shot), with low variance (0.001) and p-value < 0.01. This matches the intuition that a similar set of heads must be important within the different few-shot settings than across the zero-shot and any of the few-shot settings. However, we also see that the SRCC magnitudes for the latter are not very far off. In totality, these indicate non-trivial overlap in the (un)important attention heads for tasks across shots.

heading: 6.3 Induction Heads in OPT-66B
text: We look for induction heads in OPT-66B by quantifying the capacity of all attention heads to perform prefix matching and copying using random input sequences in a task-agnostic fashion, following the definition and algorithms by Olsson et al. (2022) discussed in §2.2 and Appendix A.8.
Figures 10a and 10b depict the prefix matching and copying score heatmaps respectively for OPT-66B. We observe that a small subset of attention heads in OPT-66B have high prefix matching scores, located in the upper layers (31+) of the model. On the other hand, there are a relatively larger number of attention heads with high copying scores, although the vast majority of these are also located in the upper layers (41+). When seen in conjunction, these observations indicate that there is a sparse set of attention heads that are capable of performing both primitive operations and thus can be deemed plausible induction heads.

heading: 6.3.1 Are Induction Heads Important?
text: We now study whether induction heads (which encode the basic in-context learning primitives of explicit prefix matching and copying) overlap with attention heads identified to be important (and consequently capable of sophisticated and latent behaviors associated with in-context learning) for our chosen downstream tasks.
A qualitative comparison of the heatmaps in Figure 10 against the heatmaps referenced in §4.1 indicates that induction heads do overlap with taskaggregated important attention heads. To better facilitate this comparison, we first formalize the total capacity of a model to perform prefix matching (or copying) to be the sum of the respective scores for individual attention heads in the model. We then investigate how much of this capacity is retained when attention heads are pruned in the order of least important heads first. Figure 11 de-
picts this comparison. We observe that much of the total prefix matching score is retained when 20% of the least important heads are removed, with the slope of decline becoming sharp only after the 40% pruning mark. This indicates that unimportant heads also have low prefix matching scores. We also observe that the prefix matching scores are generally higher for heads important for few-shot in-context learning than for heads important for zero-shot learning. On the other hand, we observe across the zero-shot and few-shot settings that the total copying score retained on pruning attention heads rapidly and consistently declines, indicating that even unimportant heads have a non-trivial capacity to perform copying. When seen in conjunction, these observations indicate that induction
heads in OPT-66B are capable of sophisticated behaviors associated with in-context learning popular downstream NLP tasks and reinforce the induction head generality arguments Olsson et al. (2022) make in the context of smaller models with stylized and synthetic tasks. We also provide per-task plots in Appendix A.9 which showcase that some tasks rely on induction heads more than other tasks.

heading: 7 Related Work
text: There has been an interest in effectively leveraging the in-context learning paradigm (Zhao et al., 2021; Holtzman et al., 2021; Min et al., 2022a; Liu et al., 2022; Lu et al., 2022; Rubin et al., 2022;
Mishra et al., 2022) 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. (2022) 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. Razeghi et al. (2022) showcase that incontext learning performance is correlated strongly with term frequencies in the pre-training corpora used. Olsson et al. (2022) 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. Chan et al. (2022) 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. Vig and Belinkov (2019) 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. Tenney et al. (2019) 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 (Voita et al., 2019), gradient-based importance and oracle knock-off importance (Michel et al., 2019), with small taskspecific trained models and report the existence of specialized heads. Given the recent trend of increasing model scale (Lieber et al., 2021; Chowdhery et al., 2022; Smith et al., 2022; Rae et al., 2021) 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 (Hoffmann et al., 2022). 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. Li et al. (2021) 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.

heading: 8 Conclusion & Future Work
text: 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 (Olsson et al., 2022) about their generality. Overall, our incontext learning-centric observations complement recent work (Hoffmann et al., 2022) 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).

heading: 9 Limitations
text: Our work is a comprehensive empirical study of a popular large language model’s capacity to perform in-context learning, relying on both task-specific (via a wide variety of challenging and practically relevant downstream tasks) and task-agnostic (via looking for induction heads) analyses and connecting the two via correlation/overlap investigations. We do not claim a causal link, i.e., we do not claim that an attention head that acquires the capacity to be an induction head will become capable of more sophisticated in-context learning associated with our downstream tasks. Making this claim will require a more deeper investigation that is outside the scope of this paper. We also do not fully understand why most attention heads seem to be unimportant for in-context learning and why there is an overlap in (un)important attention heads across tasks and shots, which warrant further investigation. Other more obvious limitations to our work include our use of only up to 5 in-context examples, random selection of in-context examples for a query input and our choice of all monolingual downstream tasks.

heading: 10 Impact Statement
text: The findings in our work have significant implications for the design, development and deployment of large language models, known to have a very high carbon footprint as well as training and inference costs. Having identified that a core nucleus of model parameters seem to be important for incontext learning, it may be possible to reduce these models’ carbon footprint and mitigate these costs. Our findings provide architectural transparency and may also be helpful in identifying targeted improvements for downstream tasks as well as for more broader facets such as bias and fairness.

heading: A Appendix
text: A.1 Head Importance Scores Figure 12 depicts the attention head aggregate importance score heatmaps in the 0-shot and 1-shot
settings. Figures 14, 15 and 16 depict the attention head importance scores for each task in the 0-shot, 1-shot and 5-shot settings respectively.
A.2 FFN Importance Scores Figure 13 depicts the task-specific and taskaveraged importance scores for feed forward networks in the 0-shot and 1-shot settings.
A.3 Removing Attention Heads Figure 17 depicts the task-specific and taskaveraged accuracy trends on iterative removal of attention heads in the order of least important first in the 0-shot and 1-shot settings.
A.4 Removing FFNs Figure 18 depicts the task-specific and taskaveraged accuracy trends on iterative removal of feed forward networks in the order of least important first in the 1-shot and 5-shot settings.
A.5 Combined Removal of Heads & FFNs Figure 19 depicts the average accuracy of all tasks on joint iterative removal of attention heads and feed forward networks in the order of least important first in the 0-shot and 1-shot settings.
A.6 Cross-Task Analysis: Spearman’s Rank Correlation
Figure 20 depicts the Spearman’s rank correlation coefficients (SRCC) between the attention head importance rankings for every pair of tasks in the 0-shot and 1-shot settings. It also depicts the SRCC between the aggregate ranking and the ranking for each constituent task.
A.7 Cross-Task Analysis: Generalization Trends
Figures 21 and 22 depict the cross-task head importance ranking generalization plots in the 0-shot and 1-shot settings.
A.8 Details of Prefix Matching and Copying Scores
Algorithms 1 and 2 contain pseudo-code to compute prefix matching and copying scores respectively for each attention head in OPT-66B. We follow the approach described by Olsson et al. (2022), but instead of computing scores using 10 sequences with fixed length of 25, we compute these scores using 100 sequences with varying lengths to account for OPT-66B’s large maximum sequence length.
As in Olsson et al. (2022), we exclude a small fraction of the most and least common tokens from the model’s vocabulary and randomly sample tokens for these sequences to strip out the effects of pretraining corpora memorization from our scores and inductive behavior analyses.
For prefix matching, the high-level approach is the following: take a random sequence, repeat it 4 times, perform a forward pass and then for each head, compute the attention pattern and take the average of all attention pattern entries attending from a given token back to tokens that succeeded the same token in earlier repeats.
For copying, the high-level approach is the following: take a random sequence, directly feed the sequence through each head and compute the contribution of the head to the output logits, and then measure how much the head increased the logit of the maximally attended to token over increasing
the logits of other attendable tokens at each timestep. Unlike Olsson et al. (2022), we do not scale the raw scores to be in the range of -1 to 1.
A.9 Importance of Induction Heads to Each Task
Figures 23 and 24 showcase the importance of induction heads to each task via measuring the percentage of the total prefix matching and copying capacities retained as a function of percentage of attention heads pruned, where heads are pruned based on each task’s head importance ranking for each in-context learning setting (zero-shot, oneshot and five-shot) in the order of least important first. A small initial slope of decline implies that unimportant heads also have low prefix matching or copying scores while a steep initial slope of decline implies unimportant heads also have high prefix matching or copying scores. We observe differ-
ences in the slopes of decline across different tasks, with tasks like HellaSwag and ReCoRD (which have high accuracies in Figure 5) having smaller initial slopes than a task like OpenBookQA (which has relatively lower accuracy in Figure 5). When seen in conjunction, these plots not only point to the generality of induction heads to more sophisticated behaviors associated with in-context learning but also indicate that some tasks rely on induction heads more than others.
Algorithm 1 Prefix Matching Scores for Attention Heads Arguments: ModelM
model← Pretrained(M) layers, heads← model.num_layers,model.num_heads_per_layer ranked_vocab_list← model.tokenizer.vocab ▷ term-frequency based vocabulary list of model exclude_vocab_size← 0.04× len(ranked_vocab_list) ▷ remove 4% most & least common tokens ranked_vocab_list← ranked_vocab_list[exclude_vocab_size : −exclude_vocab_size] prefix_matching ← [ ] for seed in {1 · · · 100} do
L← 2× seed+ 23 ▷ ensure 4L ∈ [100, 892] X ← random.choice(ranked_vocab_list, size = L, seed = seed, replace = False) ▷ L length random sequence with all unique tokens X ← repeat(X, 4) ▷ Repeat it four times Y ← model.forward(X) ▷ Forward pass the repeated sequence score← zeros(layers, heads) ▷ Zero matrix of shape layers× heads for layer in layers do
for head in heads do att← model.get_attention(layer, head) ▷ Shape: 4L× 4L for token in {L+ 1 · · · 4L} do ▷ Repetition starts from token L+ 1
att_token← att[token] ▷ Shape: 4L for every prev_token == token do ▷ Look at the previous repetitions of the token
prefix_score = att_token[prev_token+ 1] ▷ Attention given to token whose prefix is current token
score[layer][head]← score[layer][head] + prefix_score end for
end for score[layer][head]← score[layer][head]/3L ▷ Normalizing by length of for loop end for end for prefix_matching.append(score) ▷ Prefix matching scores via one randomly generated example end for prefix_matching ← average(prefix_matching) ▷ Attention head-wise average over all examples return prefix_matching
Algorithm 2 Copying Scores for Attention Heads Arguments: ModelM Definitions: Dimension per Head D,Vocabulary Size V
model← Pretrained(M) layers, heads← model.num_layers,model.num_heads_per_layer ranked_vocab_list← model.tokenizer.vocab ▷ term-frequency based vocabulary list of model exclude_vocab_size← 0.04× len(ranked_vocab_list) ▷ remove 4% most & least common tokens ranked_vocab_list← ranked_vocab_list[exclude_vocab_size : −exclude_vocab_size] copying_score← [ ] for seed in {1 · · · 100} do
L← 4× (2× seed+ 23) ▷ L ∈ [100, 892] X ← random.choice(ranked_vocab_list, size = L, seed = seed, replace = False) ▷ L length random sequence with all unique tokens score← zeros(layers, heads) ▷ Zero matrix of shape layers× heads for layer in layers do
for head in heads do attn_layer_head← model.get_attention_head(layer, head) out← attn_layer_head(X) ▷ Shape: L×D attention← model.get_attention(layer, head) ▷ Shape: L× L logits← model.hidden_to_vocab(out) ▷ Shape: L× V logits← softmax(logits, dim = 1) for token in {1 · · ·L} do
max_ind← argmax(attention[token]) ▷ Index of the token being max attended to attendable_input← X[1 : token] ▷ Attendable input tokens attendable_logits← logits[token][attendable_input] ▷ Logits of attendable tokens mean_of_logits← average(attendable_logits) raised_logits← attendable_logits−mean_of_logits relu_raised_logits← ReLU(raised_logits) ▷ Computing raise in logit values relu_raised_logit_max_ind← relu_raised_logits[X[max_ind]] temp_score← relu_raised_logit_max_ind/sum(relu_raised_logits) score[layer][head]← score[layer][head] + temp_score end for score[layer][head]← score[layer][head]/L ▷ Normalizing by length of for loop end for end for copying_score.append(score) ▷ Copying scores via one randomly generated example end for copying_score← average(copying_score) ▷ Attention head-wise average over all examples return copying_score
He lla
Sw ag PI QA
AR C
(E as
y)
AR C
(C ha
lle ng
e)
Op en
Bo ok
QA
W in
og ra
nd e
Bo ol
Q CB CO PA W IC W SC
M ul
tiR C RT E
Re Co
RD
Ag gr
eg at
e
HellaSwag PIQA
ARC (Easy) ARC (Challenge)
OpenBookQA Winogrande
BoolQ CB COPA WIC
WSC MultiRC
RTE ReCoRD Aggregate
1 0.4 0.37 0.39 0.31 0.32 0.41 0.37 0.33 0.4 0.39 0.33 0.39 0.22 0.38 0.4 1 0.38 0.37 0.31 0.32 0.4 0.36 0.36 0.36 0.37 0.34 0.36 0.22 0.37 0.37 0.38 1 0.39 0.34 0.37 0.4 0.33 0.35 0.4 0.43 0.32 0.35 0.18 0.4 0.39 0.37 0.39 1 0.33 0.33 0.41 0.36 0.35 0.41 0.4 0.34 0.37 0.19 0.39 0.31 0.31 0.34 0.33 1 0.43 0.32 0.27 0.4 0.37 0.39 0.23 0.27 0.099 0.31 0.32 0.32 0.37 0.33 0.43 1 0.34 0.28 0.44 0.39 0.41 0.22 0.28 0.13 0.32 0.41 0.4 0.4 0.41 0.32 0.34 1 0.4 0.35 0.41 0.41 0.34 0.4 0.22 0.4 0.37 0.36 0.33 0.36 0.27 0.28 0.4 1 0.29 0.37 0.36 0.39 0.42 0.26 0.37 0.33 0.36 0.35 0.35 0.4 0.44 0.35 0.29 1 0.38 0.42 0.26 0.29 0.13 0.34 0.4 0.36 0.4 0.41 0.37 0.39 0.41 0.37 0.38 1 0.44 0.32 0.36 0.18 0.38 0.39 0.37 0.43 0.4 0.39 0.41 0.41 0.36 0.42 0.44 1 0.31 0.36 0.18 0.4 0.33 0.34 0.32 0.34 0.23 0.22 0.34 0.39 0.26 0.32 0.31 1 0.38 0.26 0.35 0.39 0.36 0.35 0.37 0.27 0.28 0.4 0.42 0.29 0.36 0.36 0.38 1 0.28 0.37 0.22 0.22 0.18 0.19 0.099 0.13 0.22 0.26 0.13 0.18 0.18 0.26 0.28 1 0.22 0.38 0.37 0.4 0.39 0.31 0.32 0.4 0.37 0.34 0.38 0.4 0.35 0.37 0.22 1
Spearman Rank Correlations between Importance Score Orders
0.10
0.15
0.20
0.25
0.30
0.35
0.40
0.45
0.50
(a) 0-shot
He lla
Sw ag PI QA
AR C
(E as
y)
AR C
(C ha
lle ng
e)
Op en
Bo ok
QA
W in
og ra
nd e
Bo ol
Q CB CO PA W IC W SC
M ul
tiR C RT E
Re Co
RD
Ag gr
eg at
e
HellaSwag PIQA
ARC (Easy) ARC (Challenge)
OpenBookQA Winogrande
BoolQ CB COPA WIC
WSC MultiRC
RTE ReCoRD Aggregate
1 0.4 0.38 0.4 0.34 0.36 0.38 0.42 0.33 0.31 0.35 0.38 0.41 0.27 0.39 0.4 1 0.42 0.4 0.38 0.36 0.39 0.4 0.34 0.32 0.36 0.37 0.39 0.24 0.41 0.38 0.42 1 0.4 0.37 0.38 0.41 0.41 0.36 0.35 0.38 0.37 0.4 0.24 0.41 0.4 0.4 0.4 1 0.35 0.36 0.41 0.41 0.35 0.32 0.36 0.39 0.42 0.27 0.41 0.34 0.38 0.37 0.35 1 0.41 0.42 0.36 0.4 0.4 0.43 0.36 0.36 0.18 0.39 0.36 0.36 0.38 0.36 0.41 1 0.43 0.37 0.44 0.42 0.44 0.36 0.37 0.18 0.4 0.38 0.39 0.41 0.41 0.42 0.43 1 0.42 0.41 0.4 0.44 0.4 0.42 0.23 0.41 0.42 0.4 0.41 0.41 0.36 0.37 0.42 1 0.35 0.35 0.38 0.41 0.42 0.25 0.42 0.33 0.34 0.36 0.35 0.4 0.44 0.41 0.35 1 0.46 0.47 0.36 0.33 0.17 0.36 0.31 0.32 0.35 0.32 0.4 0.42 0.4 0.35 0.46 1 0.49 0.32 0.3 0.14 0.37 0.35 0.36 0.38 0.36 0.43 0.44 0.44 0.38 0.47 0.49 1 0.39 0.38 0.16 0.38 0.38 0.37 0.37 0.39 0.36 0.36 0.4 0.41 0.36 0.32 0.39 1 0.41 0.26 0.39 0.41 0.39 0.4 0.42 0.36 0.37 0.42 0.42 0.33 0.3 0.38 0.41 1 0.26 0.39 0.27 0.24 0.24 0.27 0.18 0.18 0.23 0.25 0.17 0.14 0.16 0.26 0.26 1 0.24 0.39 0.41 0.41 0.41 0.39 0.4 0.41 0.42 0.36 0.37 0.38 0.39 0.39 0.24 1
Spearman Rank Correlations between Importance Score Orders
0.15
0.20
0.25
0.30
0.35
0.40
0.45
0.50
(b) 1-shot
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
Section 9
A2. Did you discuss any potential risks of your work? Not applicable. Left blank.
3 A3. Do the abstract and introduction summarize the paper’s main claims? Abstract and Section 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? Section 3
3 B1. Did you cite the creators of artifacts you used? Section 3
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? The OPT-66B model is open-sourced by Meta under an unrestricted license for academic research.
3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Section 3
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Not applicable. Left blank.
B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Not applicable. Left blank.
C 3 Did you run computational experiments? Sections 3 to 6
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Section 3
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? Section 3
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Sections 3 to 6
3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Section 3
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: , 2022) 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)