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

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 11833–11856 July 9-14, 2023 ©2023 Association for Computational Linguistics Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale Hritik Bansal1∗ Karthik Gopalakrishnan2† Saket Dingliwal2 Sravan Bodapati2 Katrin Kirchhoff2 Dan Roth2 1University of California, Los Angeles 2AWS AI Labs hbansal@cs.ucla.edu {karthgop, skdin, sravanb, katrinki, drot}@amazon.com Abstract Language models have been shown to per- form better with an increase in scale on a wide variety of tasks via the in-context learn- ing paradigm. In this paper, we investigate the hypothesis that the ability of a large lan- guage model to in-context learn-perform a task is not uniformly spread across all of its un- derlying components. Using a 66 billion pa- rameter 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 hy- pothesis through a task-agnostic lens, finding that a small set of attention heads in OPT-66B score highly on their ability to perform prim- itive induction operations associated with in- context learning, namely, prefix matching and copying. These induction heads overlap with task-specific important heads, reinforcing ar- guments by Olsson et al. (2022) regarding in- duction 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. 1 Introduction 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; Chowd- hery et al., 2022; Hoffmann et al., 2022; Smith et al., 2022) based on the Transformer architec- ture (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 rev- olutionized 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 in- context learning paradigm (Brown et al., 2020). Figure 1: A sample input prompt for in-context learning and the model output. 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 sur- face 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 …

Content truncated to 500 words