Label Words as Local Task Vectors in In-Context Learning

TL;DR

This paper investigates how large language models perform in-context learning by analyzing local and global task vectors, revealing that local vectors encode rule abstractions and that global vectors may not always exist, especially in rule-dependent tasks.

Contribution

The study introduces the concept of local task vectors in ICL, showing their role in rule encoding and their convergence into global vectors in certain tasks, providing new insights into LLM mechanisms.

Findings

Local task vectors encode rule abstractions.

Global task vectors may not exist in all tasks.

ICL operates through an information aggregation mechanism.

Abstract

Large Language Models (LLMs) have demonstrated remarkable abilities, one of the most important being in-context learning (ICL). With ICL, LLMs can derive the underlying rule from a few demonstrations and provide answers that comply with the rule. Previous work hypothesized that the network creates a task vector in specific positions during ICL. The task vector can be computed by averaging across the dataset. It conveys the overall task information and can thus be considered global. Patching the global task vector allows LLMs to achieve zero-shot performance with dummy inputs comparable to few-shot learning. However, we find that such a global task vector does not exist in all tasks, especially in tasks that rely on rules that can only be inferred from multiple demonstrations, such as categorization tasks. Instead, the information provided by each demonstration is first transmitted to its answer position and forms a local task vector associated with the demonstration. In some tasks but not in categorization tasks, all demonstrations' local task vectors converge in later layers, forming the global task vector. We further show that local task vectors encode a high-level abstraction of rules extracted from the demonstrations. Our study provides novel insights into the mechanism underlying ICL in LLMs, demonstrating how ICL may be achieved through an information aggregation mechanism.