What Do Language Models Learn in Context? The Structured Task Hypothesis. (2024)

Abstract: Large language models (LLMs) exhibit an intriguing ability to learn a novel task from in-context examples presented in a demonstration, termed in-context learning (ICL). Understandably, a swath of research has been dedicated to uncovering the theories underpinning ICL. One popular hypothesis explains ICL by task selection. LLMs identify the task based on the demonstration and generalize it to the prompt. Another popular hypothesis is that ICL is a form of meta-learning, i.e., the models learn a learning algorithm at pre-training time and apply it to the demonstration. Finally, a third hypothesis argues that LLMs use the demonstration to select a composition of tasks learned during pre-training to perform ICL. In this paper, we empirically explore these three hypotheses that explain LLMs’ ability to learn in context with a suite of experiments derived from common text classification tasks. We invalidate the first two hypotheses with counterexamples and provide evidence in support of the last hypothesis. Our results suggest an LLM could learn a novel task in context via composing tasks learned during pre-training.

## introduction in-context learning (icl) is a learning paradigm where a pre-trained large language model (llm) learns to perform a certain task by extrapolating beyond a demonstration of the task in the form of example prompt-response pairs, given to the model as input. in-context learning does not require an update to the model's parameters @xcite . conditioned on the demonstration, the llm is then tasked with generating responses to additional related prompts. pre-trained large language models have exhibited an impressive ability to learn in context across various domains, e.g., code generation @xcite , education @xcite , and * equal contribution medicine @xcite . however, there is still no consensus on when or how icl works. we taxonomize existing candidate theories into three competing hypotheses (fig. 1 ), which we summarize below. hypothesis 1 (informal; task selection). during pre-training, an llm learns a set of tasks t . at inference time, the llm identifies the task τ ∈ t given the user-provided demonstration of the task and generalizes to the prompt. under hypothesis 1, the demonstration merely allows the model to recognize a task, and no actual learning takes place. @xcite offers empirical support for hypothesis 1; they show that randomly shuffling the responses in the demonstration hardly has any effect on icl performance, suggesting the demonstration of the task only serves to enable the llm to look up a task. in other words, hypothesis 1 asserts that no learning takes place during icl. some authors @xcite @xcite have also argued for hypothesis 1 from a theoretical angle, contending that if an llm is pre-trained on a corpus that is generated from a mixture model over tasks, it will be able to infer the task that generated the demonstrations to be able to generalize to a prompt not in the demonstration. the next hypothesis revolves around metalearning @xcite @xcite . hypothesis 2 (informal; meta-learning). during pre-training, an llm learns certain learning algorithms. during icl, the llm learns a task τ directly from the demonstration using one of the learned learning algorithms. hypothesis 2 suggests that the pre-training stage prepares the parameters in an llm in such a way that various learning algorithms, e.g., gradient descent and least squares regression, can be implicitly deployed during icl @xcite @xcite to learn a new task from a demonstration. however, the setting assumed in the above-cited theoretical development is usually over-simplified, e.g., the assumption that the attention mechanism is linear. moreover, the empirical evidence for hypothesis 2 is mostly derived from experiments on small transformers trained from scratch on synthetic data @xcite @xcite . the final hypothesis may be viewed as a mixture of hypothesis 1 and hypothesis 2. hypothesis 3 (informal; structured task selection). during pre-training, an llm learns a set of tasks t . at inference time, the llm uses the demonstration to compose a sequence of learned tasks τ 1 , τ 2 , . . . ∈ t and uses this composition for prediction. the composition itself may result in a novel task not seen during pre-training. hypothesis 3 extends hypothesis 1 by allowing icl to not only index into the pre-learned task set t , but also compose tasks to obtain a novel task τ = τ 1 • τ 2 • • • • / ∈ t . 1 hahn and goyal (2023) lay the theoretical groundwork for hypothesis 3; under their framework, they argue that the compositional structure of linguistic pre-training data gives rise to the task composition ability of icl. in this paper, we conduct a comprehensive battery of experiments to examine the above-stated three hypotheses. given a demonstration of a task, i.e., a sequence of prompt-response pairs from which an llm can successfully learn the task in context, we create a response-altered (ra) task by altering the responses such that the new task is unlikely to have occurred during pre-training. our results confirm that llms can learn such an ra task in context, which rejects hypothesis 1. we then create prompt-altered (pa) tasks, which instead alter the prompts. if hypothesis 2 is true, this modification should not affect the performance of icl. however, we find that the llms yield a substantially worse performance on pa tasks than on ra tasks, which contradicts hypothesis 2. lastly, in support of hypothesis 3, we identify a sequence of simple tasks that the llms can compose 1 the definition of task composition • will be given in § 2.2. to obtain unobserved ra tasks. it offers a possible explanation for icl's ability to perform novel tasks. ## testing hypothesis 1 according to hypothesis 1, icl selects the task it needs to perform from the demonstration. however, it is only able to select among the finite set of tasks observed at the pre-training time, denoted as t . it follows from hypothesis 1, then, that if a novel task that has never been seen during pretraining is presented to a pre-trained model as a demonstration, icl should not be able to perform it. we construct such a novel task as follows. we first create a string-to-string function g : σ * → σ * . note that such a function g is a special case of a response distribution that places probability 1 on a specific output string for every input string. then, given a task ⟨π τ , ρ τ ⟩, we obtain a new task τ ra by applying g to the responses of τ , i.e., note that the definition in ( 5 ) is no more than task composition, i.e., ⟨g, •⟩ • ⟨ρ, π⟩ where • is a stand-in for an arbitrary prompt distribution. we define τ g @xmath0, •⟩ for the remainder of the paper, i.e., a task induced by the string-to-string function g with a stand-in prompt distribution; we also write τ @xmath1. the new task τ ra is almost certainly not observed in the pre-training data, i.e., τ ra / ∈ t . we call τ ra a response-altered task (ra) and the icl setting with ra tasks τ ra -icl. this setting resembles the semantically unrelated label icl setting of @xcite and the abstract formalization of @xcite . we examine the following logical consequence of hypothesis 1. prediction 1. if hypothesis 1 is true, then an llm's performance in the τ ra -icl setting should be similar to random guessing. ## testing hypothesis 2 we have shown in the previous section that an llm can learn a novel task in context, but it is still unclear what type of tasks can be learned in context. if we accept hypothesis 2 as true, certain learning algorithms are learned by the pre-trained llm, so learning from a demonstration in context happens on the fly at inference time. this hypothesis implies that there need not be knowledge of the prompted task in the pre-training data. and indeed, on this view, the role of the pre-training is merely to prepare the parameters in the llm in such a way that the llm architecture encodes various learning algorithms. for instance, some authors @xcite @xcite argue that training a linear model with gradient descent can be encoded as in-context learning under certain simplifying assumptions. this leads to the following prediction. prediction 2. if hypothesis 2 is true, then icl should behave similarly to a model trained with a certain learning algorithm, e.g., gradient descent. a caveat of prediction 2 is, of course, that we do not know a priori which learning algorithm the llms learn at pre-training time. ## testing hypothesis 3 because neither hypothesis 1 nor hypothesis 2 matches our experimental findings, we now turn to hypothesis 3 to explain the empirical facts. ## conclusion in this paper, we examine the ability of pre-trained llms to learn various tasks in context. we find that llms can learn text classification tasks with corrupted responses (τ ra -icl), but when the prompts are corrupted in the same manner (τ pa -icl), icl performs significantly worse than a logistic regression model. a closer look suggests the tasks in the τ ra -icl may have been learned via composing primitive tasks learned during pre-training. overall, our paper provides insights into the nature, abilities, and limitations of icl.