[Paper Review] In-Context Learning Creates Task Vectors

Roee Hendel, Mor Geva+1 more|arXiv (Cornell University)|Oct 24, 2023

Topic Modeling4 citations

TL;DR

This paper proposes that in-context learning (ICL) in large language models (LLMs) operates by compressing demonstrations into a single task vector, which then modulates the model’s output for a query. The authors demonstrate that ICL can be viewed as a hypothesis class learning mechanism where the task vector captures the task rule, enabling accurate predictions without direct access to the demonstrations, with empirical validation across multiple models and tasks.

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θ(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.

Motivation & Objective

• To understand the underlying mechanism of in-context learning (ICL) in large language models (LLMs), which remains poorly understood despite its empirical success.
• To map ICL to a formal learning-theoretic framework by identifying a hypothesis class where ICL operates.
• To investigate whether ICL can be decomposed into a two-stage process: learning a task vector from demonstrations, followed by applying it to a query.
• To validate that the task vector captures meaningful, task-specific semantic information independent of the input demonstrations.
• To explore the implications of this mechanism for efficient and interpretable LLM adaptation to new tasks.

Proposed method

• Propose a hypothesis class view of ICL, where the model learns a task vector θ(S) from demonstrations S, independent of the query x.
• Define a two-part mechanism: (1) a 'learning algorithm' 𝒜 that computes θ(S) from S, and (2) a 'rule application' f that uses θ(S) and x to generate the output, without direct access to S.
• Use attention mechanisms in decoder-only transformers to extract the task vector as a hidden state representation after processing S and x.
• Project the task vector into the vocabulary space via a linear projection to interpret its semantic content.
• Train and evaluate the method on diverse NLP tasks, including translation, word sense disambiguation, and arithmetic reasoning.
• Compare performance using the task vector-based method against standard ICL and baseline models to validate effectiveness and generalization.

Figure 1: ICL as learning in a Hypothesis Class. 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: $\mathcal{A}$ , a “le

Experimental results

Research questions

• RQ1 Can in-context learning in LLMs be formalized as a hypothesis class learning mechanism, where a single task vector θ(S) encodes the rule from demonstrations S?
• RQ2 Is the task vector θ(S) independent of the query x, such that it can be applied to any new query without reprocessing S?
• RQ3 Does the task vector θ(S) contain meaningful, interpretable semantic information about the task, even when the task-specific terms do not appear in the input?
• RQ4 Can the task vector approach generalize across different models and tasks while maintaining high accuracy?
• RQ5 How does the performance of the task vector-based method compare to standard ICL and baseline models across diverse tasks?

Key findings

• The task vector θ(S) successfully captures the underlying rule from demonstrations S, enabling accurate prediction on unseen queries without direct access to S.
• On the translation task, the model achieved 96% accuracy using the task vector, demonstrating strong generalization from few-shot examples.
• The top tokens in the task vector projection for the French-to-English translation task included 'English' and 'translate', indicating that the vector encodes semantic task information not present in the input.
• For the 'Present to Past' task, the model achieved 95% accuracy, showing robust performance across diverse linguistic transformations.
• The method generalizes across models, with strong results observed on LLaMA 13B, GPT-3, and other models across multiple tasks.
• The task vector approach provides a mechanistic interpretation of ICL, suggesting that LLMs compress demonstrations into a single, interpretable vector that modulates the model's behavior.

Figure 2: Separating $\mathcal{A}$ and $f$ . To make $\boldsymbol{\theta}$ independent of the query $x$ , we use a dummy query ( $x^{\prime}=\mbox{Plum}$ ) and use the representation of $\rightarrow$ at the $L^{th}$ layer as $\boldsymbol{\theta}$ . The vector $\boldsymbol{\theta}$ is then patched at