[Paper Review] Function Vectors in Large Language Models

Eric Todd, Millicent Li+4 more|arXiv (Cornell University)|Oct 23, 2023

Topic Modeling4 citations

TL;DR

This paper introduces function vectors (FVs)—compact, task-specific representations in autoregressive transformers that encode input-output functions and can be transferred across diverse contexts. Using causal mediation analysis, the authors identify attention heads that transport FVs, which robustly trigger task execution even in zero-shot and natural language settings, and demonstrate that FVs support compositional vector algebra for complex function synthesis beyond word-level semantics.

ABSTRACT

We report the presence of a simple neural mechanism that represents an input-output function as a vector within autoregressive transformer language models (LMs). Using causal mediation analysis on a diverse range of in-context-learning (ICL) tasks, we find that a small number attention heads transport a compact representation of the demonstrated task, which we call a function vector (FV). FVs are robust to changes in context, i.e., they trigger execution of the task on inputs such as zero-shot and natural text settings that do not resemble the ICL contexts from which they are collected. We test FVs across a range of tasks, models, and layers and find strong causal effects across settings in middle layers. We investigate the internal structure of FVs and find while that they often contain information that encodes the output space of the function, this information alone is not sufficient to reconstruct an FV. Finally, we test semantic vector composition in FVs, and find that to some extent they can be summed to create vectors that trigger new complex tasks. Our findings show that compact, causal internal vector representations of function abstractions can be explicitly extracted from LLMs. Our code and data are available at https://functions.baulab.info.

Motivation & Objective

• To investigate whether large language models (LLMs) contain internal, reusable representations of input-output functions.
• To determine if these representations, termed function vectors (FVs), can be extracted and transferred across different input contexts, including zero-shot and natural language settings.
• To examine whether FVs encode more than just output vocabulary, and whether they support compositional vector operations for creating new functions.
• To evaluate the robustness and generalization of FVs across model architectures, layers, and task types.

Proposed method

• Causal mediation analysis is used to identify attention heads that mediate information flow during in-context learning (ICL), isolating the contribution of specific heads to task execution.
• Function vectors are constructed by summing activations from identified attention heads, forming a compact representation of the demonstrated task.
• Activation patching is applied by inserting the FV into a new context at a specific layer, measuring its effect on zero-shot task performance.
• Decoding analysis evaluates whether FVs encode the output space of the function, using top-k token probabilities from the model's next-token prediction.
• Vector composition experiments test whether summing FVs of simpler tasks produces new FVs capable of executing combined, complex tasks.
• Experiments are conducted across multiple models (GPT-J, Llama 2, etc.), varying from 6B to 70B parameters, and across diverse ICL tasks (e.g., translation, antonym generation, sentiment).

Experimental results

Research questions

• RQ1 Do large language models contain compact, reusable internal representations of input-output functions that can be transferred across different contexts?
• RQ2 Can function vectors (FVs) trigger task execution in zero-shot and natural language settings, even when the input context bears no resemblance to the original ICL demonstration?
• RQ3 To what extent do FVs encode the output space of a function, and is this encoding sufficient to reconstruct the FV?
• RQ4 Can FVs be composed via vector addition to form new FVs that trigger complex, composite tasks not directly demonstrated?
• RQ5 Are FVs robust across model sizes and layers, and do they consistently improve zero-shot performance in middle layers?

Key findings

• Function vectors (FVs) are robustly identified in middle layers of transformer models and can trigger task execution in zero-shot and natural language contexts, with up to 24.3% accuracy for antonym generation in GPT-J.
• FVs extracted from one context can successfully induce task execution in unrelated contexts, including those with no structural similarity to the original ICL examples.
• Although FVs often encode the output vocabulary of a function, this information alone is insufficient to reconstruct the FV, indicating additional structural or semantic content is encoded.
• Vector composition of FVs enables the creation of new FVs that trigger complex, composite tasks—for example, summing FVs for 'capitalize first letter' and 'antonym' can produce a vector that performs both operations.
• FV performance is most effective when inserted into early to middle layers, and this pattern holds consistently across model sizes from 6B to 70B parameters (e.g., Llama 2).
• In Llama 2 (13B), vector-composed FVs achieved 87% accuracy on a combined 'capitalize first letter' task, outperforming individual FVs and demonstrating compositional generalization.