Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations

TL;DR

This paper investigates the emergence and functionality of task vectors in in-context learning, proposing a linear combination hypothesis supported by theoretical and empirical evidence, and exploring their limitations and potential enhancements.

Contribution

It introduces the Linear Combination Conjecture for task vectors, supported by analysis and experiments, advancing understanding of their role in in-context learning.

Findings

Task vectors naturally emerge in linear transformers trained on triplet prompts.

Task vectors fail to represent high-rank mappings, confirmed on practical LLMs.

Injecting multiple task vectors into prompts can enhance in-context learning performance.

Abstract

Task vectors offer a compelling mechanism for accelerating inference in in-context learning (ICL) by distilling task-specific information into a single, reusable representation. Despite their empirical success, the underlying principles governing their emergence and functionality remain unclear. This work proposes the Linear Combination Conjecture, positing that task vectors act as single in-context demonstrations formed through linear combinations of the original ones. We provide both theoretical and empirical support for this conjecture. First, we show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis. Next, we predict the failure of task vectors on representing high-rank mappings and confirm this on practical LLMs. Our findings are further validated through saliency analyses and parameter visualization, suggesting an enhancement of task vectors by injecting multiple ones into few-shot prompts. Together, our results advance the understanding of task vectors and shed light on the mechanisms underlying ICL in transformer-based models.