Source-linked AI summary
In-Context Learning Creates Task Vectors
Roee Hendel, Mor Geva, Amir Globerson
TL;DR
The paper addresses the unclear relationship between ICL and standard hypothesis-class learning. It proposes that demonstrations are compressed into a task vector used by a transformer on the query, and reports that this decomposition approximates ICL across models and tasks.
Problem
ICL processes demonstrations and queries jointly through a transformer, so its underlying hypothesis space and learning mechanism remain unclear.
Method
The authors decompose ICL into a learning algorithm that computes a task vector from demonstrations and a rule that applies it to the query without direct access to the demonstrations.
Results
The proposed separation retains around 80-90% of regular ICL accuracy across models, while the baseline reaches only 10-20%.
Takeaways & Limitations
ICL can be viewed as compressing demonstrations into a single task vector that guides the transformer in generating outputs for queries.
Takeaways & Limitations
The study focuses on cases where a single task vector suffices and outputs are single tokens; more complex tasks may require more elaborate parameterization and multi-token outputs.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper asks how ICL uses demonstrations to produce predictions and argues that it can be understood through a hypothesis-class view. Across models and diverse tasks, the authors propose that demonstrations are compressed into a task vector that modulates the transformer.
- ICL lets LLMs learn new rules from very few demonstrations, such as mapping “Corn” to “Yellow” after seeing two examples.
- Standard hypothesis-class learning represents functions as h(x; θ), but ICL directly processes the concatenation [S, x], leaving its hypothesis space unclear.
- The paper proposes that ICL maps demonstrations S into a task vector θ(S) representing the demonstrated rule, then applies a function f(x; θ).
- Task vectors resemble soft prompts because both modulate the transformer toward a task, but ICL computes them during the forward pass rather than fine-tuning them.
- The authors validate this hypothesis-class view through experiments on publicly available LLMs and a diverse set of tasks.
2 A Hypothesis Class View of ICL
The paper formulates ICL as a hypothesis class in which a learning algorithm maps demonstrations to a query-independent task vector, and a separate rule applies that vector to the query. It proposes a transformer-based procedure to isolate these components.
- 2 A Hypothesis Class View of ICL: The hypothesis-class view asks whether ICL converts demonstrations S into θ, the parameter of a function operating on query x.
- 2 A Hypothesis Class View of ICL: The framework decomposes ICL into A, which maps S to a query-independent task vector θ, and f, which maps x to an output using θ without direct access to S.
- 2 A Hypothesis Class View of ICL: This decomposition defines the hypothesis class H = {f(·; θ) | θ}.
- 2 A Hypothesis Class View of ICL: The proposed realization treats early transformer layers as computing θ from S and later layers as applying the resulting rule to x.
- 2.2 A Proposed Hypothesis Class: Experiments aim to isolate A and f and test whether θ vectors are interpretable and correspond to learned tasks.
- 2.2 A Proposed Hypothesis Class: A dummy query x′ and activation patching are used to make θ independent of x and prevent f from directly accessing S.
3 Validity of the Hypothesis Class View
The authors isolate ICL into task-vector calculation and query-time rule application, then test whether this separation preserves performance across models and tasks. The separated procedure retains much of regular ICL’s accuracy while substantially outperforming the no-demonstration baseline.
- 3.1 Separating A and f: The procedure separates ICL into A, which computes θ, and f, which applies θ to query x without direct access to S.A dummy query x′ helps prevent θ from depending on x, while patching θ into a pass over x prevents f from accessing S.
- 3.2 Experimental Setup: The study evaluates 18 tasks across algorithmic, translation, linguistic, and factual-knowledge categories, using single-token outputs and multiple open LLMs.The models include LLaMA, GPT-J, and Pythia variants.
- 3.3 Finding L: All models exhibit a performance peak at a similar intermediate layer L, despite differences in parameter counts and layer numbers.The layer is selected by evaluating development-set accuracy across different choices of L.
- 3.4 Accuracy of Hypothesis Based Prediction: Across all models, the Hypothesis procedure retains around 80-90% of regular ICL accuracy, whereas the Baseline reaches only 10-20%.The comparison uses average accuracy across all tasks for Regular, Hypothesis, and Baseline procedures.
4 Robustness of Task Vectors
The task vector remains stable when the demonstrations and dummy query vary within a task. Across tasks, vectors form distinct clusters, with same-category tasks also appearing closer.
- Robustness analyses: Task vectors are generated from 50 varied choices of S and x′ per task using LLaMA 7B for robustness analyses.The analyses examine both the geometry and variability of θ.
- Geometry of θ: Task vectors form distinct t-SNE clusters, each containing vectors from a single task, while tasks in the same category show proximity.The clustering supports the interpretation that vectors encapsulate task understanding.
- Variability of θ: Within-task vector distances are smaller than across-task distances, indicating stability within tasks and limited influence from x′ or S.This comparison is shown through histograms of distances within and across tasks.
5 Dominance of θ Patching
A conflicting-task experiment tests whether the query-time computation follows the injected task vector rather than the demonstrations. The results show that patching θ shifts predictions toward the task encoded by that vector.
- Conflicting Tasks Experiment: The experiment compares a Regular pass using demonstrations S_A with a Conflicting pass that retains S_A while injecting θ_B.Tasks A and B share the input space but differ in their outputs.
- Results: Regular inference achieves 90%+ accuracy on task A, while the Conflicting pass achieves high accuracy on task B despite retaining demonstrations from task A.Task B accuracy is slightly lower, potentially because of the performance dip and the presence of S_A.
- Interpretation: The conflicting-task results imply that the model mainly relies on θ and largely disregards demonstrations S_A during the patched computation.The authors qualify that task B accuracy is slightly low.
6 Interpreting θ
The authors probe the semantic content of θ by projecting its hidden state into vocabulary-token distributions. The resulting top tokens often directly describe the demonstrated task, even when those terms were absent from the context.
- Method: Because θ is an intermediate hidden state, the authors use vocabulary projection to inspect the top tokens induced by that state.This provides an interpretability analysis of the information encoded in θ.
- Results: For several tasks, top tokens directly describe the task; French-to-English translation, for example, produces tokens such as “English” and “translate”.These terms never explicitly appeared in the context.
- Interpretation: The observed task-descriptive tokens support the view that θ carries significant, non-trivial semantic information about the task.The evidence comes from the vocabulary distributions induced by θ.
7 Related Work
Prior work studies how in-context learning emerges and how transformers can implement learning algorithms, while this work identifies a parameter space for LLMs.
- ICL in LLMs: Studies of ICL emergence examine how pre-training data structure and distributional properties affect the capability.The cited work includes probabilistic analyses using Hidden Markov Models and empirical studies of distributional properties.
- Meta-Learning in Transformers: Meta-learning studies typically train transformers from scratch on elementary tasks and draw theoretical parallels with algorithms such as Gradient Descent.These studies often focus on tasks such as linear regression and assume a known parameter space for gradient descent.
- ICL in LLMs: Related LLM research connects ICL to induction heads and Gradient Descent, while concurrent work studies single vectors that encode learned functions.The paper characterizes its findings as complementary to the concurrent task-vector work.
8 Conclusions
The paper presents ICL as compressing demonstrations into a single task vector that guides the transformer’s outputs for queries. It identifies understanding task-vector construction and use as directions for future work.
- Conclusions: ICL compresses a training set into a single task vector that guides the transformer to generate outputs for queries.The task vector provides a simple structure for understanding the learning mechanism described in the paper.
- Conclusions: Future work should investigate how the task vector is constructed and how it is used to calculate the output.
Limitations
The study evaluates relatively simple, single-token tasks, so it remains uncertain whether the observed task-vector mechanisms extend to more complex ICL settings.
- Scope: The study focuses on relatively simple tasks, whereas ICL can also perform more complex tasks such as arithmetic reasoning.Whether the observed mechanisms translate to those cases remains unresolved.
- Scope: The approach targets cases where a single task vector suffices, while more complex ICL cases may require more elaborate parameterization.
- Scope: The experiments focus on tasks with single-token outputs, although some tasks require multi-token outputs.
- Mechanistic understanding: The paper does not provide a mechanistic explanation of how the task vector is formed or used by the transformer’s parameters.
A Appendix
The appendix provides additional details and results.
- Appendix: The appendix provides additional details.
- Appendix: The appendix provides additional results.
- Appendix: The appendix material supplements the main paper’s presentation.
A.1 Additional Details
The appendix documents the study’s 18 tasks, models, data sources, experiments, and expanded analyses. Additional results examine layer choice, task-vector geometry, vocabulary projections, and task- and model-specific performance.
- Tasks: The study covers 18 tasks across algorithmic, translation, linguistic, and knowledge categories, with single-token outputs.The appendix provides complete task descriptions and a representative task table.
- Task Data: Task data come from programmatically generated algorithmic examples, multilingual word lists, linguistic datasets, and a counterfactual knowledge dataset.The sources differ by task category and are documented alongside repository materials.
- Layer Selection: Accuracy peaks at a similar intermediate layer across models despite differences in parameter counts and layer numbers.The result is reported for LLaMA models, with additional models discussed in the appendix.
- Vocabulary Projections: Expanded vocabulary projections provide top-token analyses for additional tasks and multiple LLMs.The appendix extends the three-task LLaMA 13B inspection presented in the main text.
- Conflicting Tasks: The conflicting-tasks experiment injects a task vector from Task B while demonstrations specify Task A.The regular condition uses only Task A demonstrations, whereas the conflicting condition combines Task A demonstrations with θ(SB).
- Task-Vector Geometry: Within-task task vectors are closer than across-task vectors, indicating stability across different training sets and dummy queries.Figure 8 compares within-task and cross-task distance distributions.