In-Context Learning Creates Task Vectors (2023)

Abstract: In-context learning (ICL) in Large Language Models (LLMs) has emerged as a powerful new learning paradigm. However, its underlying mechanism is still not well understood. In particular, it is challenging to map it to the “standard’ machine learning framework, where one uses a training set S to find a best-fitting function f(x) in some hypothesis class. Here we make progress on this problem by showing that the functions learned by ICL often have a very simple structure: they correspond to the transformer LLM whose only inputs are the query x and a single “task vector’ calculated from the training set. Thus, ICL can be seen as compressing S into a single task vector \boldsymbol{\theta}(S) and then using this task vector to modulate the transformer to produce the output. We support the above claim via comprehensive experiments across a range of models and tasks.

## introduction large language models have improved dramatically over the last several years. one striking property of these models is that they can learn new rules from very few demonstrations. for instance, a model can be prompted with the input "@xmath0, @xmath1, @xmath2. the model has thus learned a mapping based on just two examples, which it can apply correctly to new examples. this capability, referred to as in-context learning (icl), has been used extensively, yielding impressive empirical results @xcite @xcite . given this success, it is natural to ask what is the underlying mechanism behind icl. namely, how does the model internally use the demonstrations s and the query x to produce the required output? here we approach this question by utilizing the 1 we release our code at https://github.com/ roeehendel/icl_task_vectors . in icl, one provides an llm with a prompt including demonstrations s of some task, and a query x. the model generates the output for x (here "yellow"). we show that the underlying process can be broken down into two parts: a, a "learning algorithm" (marked in blue), computes a query-agnostic vector θ(s), which we view as a parameter of a function in a hypothesis class. the second part, denoted by f and marked in yellow, is the application of the rule defined by θ on the query x, without direct dependence on s. it is unclear whether icl operates in such a way because the prediction is performed via t ([s, x]), where t is typically an auto-regressive transformer and [s, x] is a concatenation of the tokens in s and x. thus, in the general case, it can be an arbitrary function that operates on s and x to produce the output. this can include "non-parametric" methods such as nearest-neighbor. recent work has begun to explore this question. for example, it was shown that when training a transformer from scratch to perform linear regression in context, the emerging learning algorithm is similar to stochastic gradient descent @xcite . however, for llms performing more complex natural language tasks, it is not at all clear what the hypothesis space may be. in this work, we show that on a wide range of tasks, icl in llms can be viewed as working on a very natural hypothesis space. we argue that, given a training set s, the transformer maps it into a "task vector" θ(s) that essentially represents the mapping/rule described in s. foot_0 namely, given the transformer t and a vector θ, we can construct a new function f (x; θ) that implements the task. the function f is very similar to the original transformer applied to x without demonstrations but instead modulated by θ (see fig. 2 ). our view is also related to soft prompts @xcite , since both approaches modulate the function of the transformer towards a particular task. however, in icl, task vectors are calculated in the forward pass rather than being fine-tuned. our contributions include proposing a hypothesis-class based mechanistic view of icl, and conducting experiments to validate our view on a range of publicly available llms and a diverse set of tasks. our results further the understanding of icl and may have practical implications for the efficient adaptation of llms to perform specific tasks. ## a hypothesis class view of icl motivated by the hypothesis class view of learning theory, our goal is to understand if icl maps the set of demonstrations s to a function on the query x and how this mapping occurs. specifically, we seek to see if icl converts s into θ -the "parameters" of a function within a certain hypothesis space. our empirical findings suggest this view is applicable, shedding light on the structure of the hypothesis space on which icl can be viewed to operate. ## validity of the hypothesis class view we first show that separating the forward pass into the two distinct components a and f , defined in §2.2, maintains the high accuracy of icl. ## robustness of task vectors in our setting, θ is derived from s and a dummy query x ′ . it is natural to examine the robustness of θ to variations in these inputs. intuitively, if it represents the task, it should remain stable across different s and x ′ values. to test this, we use llama 7b to generate 50 task vectors per task with varied s and x ′ and conduct two analyses. geometry of θ a t-sne dimensionality reduction (fig. 5 ) reveals that the task vectors form distinct clusters, each containing task vectors of a single task. fig. 9 further shows proximity between tasks of the same category, strengthening the idea that they encapsulate task understanding. variability of θ fig. 8 shows histograms of distances within and across tasks. it can be seen that vectors within the same task are closer than those between different tasks, indicating that θ is stable within tasks and not highly influenced by x ′ or s. ## dominance of θ patching in §3 we prevented f from directly accessing s. however, in a regular forward pass during icl, the last token can attend to s. here we verify that even in this case, f mainly uses the task vector θ, without directly accessing the demonstrations s. to this end, we use a pair of tasks, a and b, sharing the input space but differing on the output. we first use a "regular" forward pass, where we provide the model with demonstrations s for task a (denoted s a ), to verify the model can perform this task using icl. then, we do a "conflicting" forward pass, still providing s a , while injecting θ b . for more details, refer to fig. 6 in tab.2, the "regular" forward pass shows high accuracy on task a (90%+), as anticipated. however, the "conflicting" forward pass yields high accuracy on task b, corresponding to the injected task vector θ. this implies that the model mainly relies on θ, largely disregarding the demonstrations s for task a. we note that the accuracy on task b is slightly low, likely consistent with the performance dip seen in fig. 6 , and potentially further affected by the presence of s. ## interpreting θ the learned vector θ intuitively captures information about the task demonstrated by s. here we provide evidence supporting this interpretation. since θ is an intermediate hidden state of the transformer, we can employ a vocabulary projection method @xcite . namely, we examine the top tokens in the distribution over the vocabulary induced by the hidden state. tab. 3 shows the top tokens for three tasks for llama 13b (more models and tasks are provided in tab. 7 in §a). in multiple cases, we observe tokens that directly describe the task. importantly, these terms never explicitly appeared in the context. for example in the task of translation from french to english, we observe tokens such as "english" and "translate". this supports our view that θ carries significant, non-trivial semantic information about the task. ## related work emergence of icl a key question with icl is how it emerges as a capability from pre-training the llms. @xcite provides results in this direction that highlight the importance of training data structure. xie et al. use probabilistic analysis and model pre-training data using hidden markov models to theoretically explain the emergence of icl, while @xcite empirically explore the effect of several distributional properties of the pre-training data. ## conclusions through this exploration of icl in llms, we have shed light on a new perspective of icl learning mechanisms. we have revealed a simple and elegant structure: icl functions by compressing a given training set into a single task vector, which then guides the transformer to generate appropriate outputs given queries. our work provides a stepping stone towards understanding how llms perform icl. in light of our findings, future work could focus on understanding how the task vector is constructed as well as how it is used to calculate the output.