PaperlayerView exact revision on arXiv

Source-linked AI summary

Function Vectors in Large Language Models

Eric Todd, Millicent L. Li, Arnab Sen Sharma, Aaron Mueller, Byron C. Wallace, David Bau

arXiv:2310.15213v2cs.CLcs.LG

TL;DR

The paper asks how language models represent and execute complex in-context functions. Using causal mediation across diverse tasks, it extracts function vectors from a small set of attention heads and tests their portability, internal information, and composition. The results show that these vectors are robust task representations that can sometimes combine to execute new complex functions, while their output vocabulary alone is insufficient to reconstruct them.

  • Problem

    How autoregressive transformers represent and execute complex in-context-learning functions remains incompletely understood.

  • Method

    The paper uses causal mediation analysis to identify attention heads transporting task representations and extracts function vectors by summing their outputs.

  • Results

    Function vectors are robust across unrelated contexts, often encode output vocabulary without being reducible to it, and can sometimes compose into vectors that execute complex tasks.

  • Takeaways & Limitations

    The findings provide evidence that language models contain compact, causal internal vector representations of function abstractions that can be explicitly extracted.

  • Takeaways & Limitations

    Function vectors are not a complete account of in-context learning, and greater transparency could also enable abuse such as injecting or amplifying undesirable functions.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper identifies function vectors as compact hidden-state representations of in-context tasks and shows that they can trigger task execution across contexts. These vectors are robust, contain more than output-vocabulary information, and can sometimes compose into vectors for complex tasks.

  • Core contribution: Function vectors are compact vector representations of input-output tasks found in transformer hidden states during in-context learning.They trigger execution of a procedure rather than directly performing the task.
  • Core contribution: Causal mediation analysis identifies a small set of attention heads whose summed outputs transport the task representation.The study evaluates over 40 diverse ICL tasks across language models ranging from 6B to 70B parameters.
  • Portability: Function vectors typically trigger task execution in differently formatted, zero-shot, and natural-text contexts unlike the contexts from which they were extracted.This portability is tested across diverse input settings.
  • Internal structure: Function vectors often encode a function’s output vocabulary, but that vocabulary alone does not fully identify the vector.Additional information is required to invoke the function.
  • Composition: To some extent, summed function vectors produce new vectors that execute complex tasks combining their constituent tasks.The resulting function-vector algebra is distinct from word-embedding arithmetic and can specify nonlinear tasks such as antonym generation.

2 METHOD

The method extracts task representations by tracing causal information flow from successful ICL examples through attention heads, then tests the resulting vectors by intervening on hidden states. Heads are selected by their indirect effects on recovering correct answers across tasks.

  • Motivating observation: The motivating test averages last-token activations from ICL prompts and adds them to an unseen zero-shot prompt at a chosen layer.For GPT-J antonym generation, adding the average activation at layer 12 yields 24.3% accuracy.
  • Prompt construction: The analysis uses successful ICL prompts containing task-demonstrating input-output pairs and a novel query whose correct answer is ranked highest by the model.Uninformative prompts pair inputs with randomly shuffled outputs to disrupt the demonstrated relationship.
  • Causal mediation: Causal mediation replaces an attention head’s activation on an uninformative prompt with its mean task-conditioned activation and measures the change in probability assigned to the correct answer.The average indirect effect aggregates this causal difference across tasks and corrupted prompts.
  • Head identification: The most influential heads are concentrated in early-middle or middle layers and attend most strongly to output tokens of the ICL exemplars.The top 10 heads are selected by average indirect effect across tasks.
  • Function-vector extraction: A function vector is formed by summing the average outputs of the selected causal attention heads for a task.Its causal effect is tested by adding it to hidden states as the model resolves a prompt and measuring task execution.

3 EXPERIMENTS

The experiments test function vectors across models, tasks, layers, contexts, and compositions. FVs improve task execution in unfamiliar contexts, encode more than output vocabularies, and can sometimes compose into new functions.

  • FV evaluation: Adding FVs improves performance in shuffled-label and zero-shot contexts, outperforming layer-averaged vectors for GPT-J.Table 2 averages accuracy across six tasks and random seeds at approximately one-third network depth.
  • Layer effects: FV causal effects are strongest in early-middle layers and drop sharply to near-zero in late layers across tasks, architectures, and model sizes.The layer pattern suggests FVs trigger later nonlinear computations rather than acting as simple linear offsets.
  • Portability: 76.2±13.8% shuffled-label accuracy and 40.0±16.7% zero-shot accuracy with FVs exceed GPT-J baselines of 32.3 ± 12.8% and 6.2 ± 4.3%.These results average across six tasks and 20 prompt templates.
  • Portability: FVs retain task-triggering effects across differently formatted, zero-shot, and natural-text contexts, including completions beyond the original demonstrated task.Natural-text results for all six tasks are consistent with portability from ICL contexts.
  • FV structure: Decoded FV tokens often reflect a task’s output space, but reconstructing a working FV usually requires more than 100 vocabulary tokens.Thus, top decoded tokens do not fully identify the function vector.
  • Vector algebra: Some algebraic FV compositions outperform individual FVs and ICL on complex tasks, while other tasks resist composition.The demonstrated compositions combine constituent word-selection and word-transformation behaviors.

4 RELATED WORK

The paper positions function vectors within research on task representations, in-context learning, mechanistic interpretability, attention, and activation steering. Its distinction is that FVs are compact, causally identified representations already present in LLM hidden states.

  • Task Representations: Unlike prior task encodings based on metamappings, codebooks, soft prompts, or parameter perturbations, FVs already exist within LLMs and can be extracted.The paper also distinguishes FVs from hidden-state task clusters because FVs have causal effects.
  • In-Context Learning: The work addresses unresolved mechanisms of complex ICL execution beyond prior accounts emphasizing copying, theoretical inference, prompt form, or Bayesian task inference.It characterizes function vectors as compact representations of input-output tasks.
  • Mechanistic Interpretability: The analysis builds on mechanistic-interpretability methods for causal mediation and uses the logit lens to examine FV vocabulary.It is related to studies of ICL copying and components involved in task execution.
  • Analyzing the Attention Mechanism: Unlike attention-weight analyses, the paper focuses on transported attention-head outputs because attention weights alone do not fully explain model outputs.The emphasis is on content carried by attention-mediated representations.

5 DISCUSSION

The discussion presents function vectors as a simple, shared mechanism for task representation in transformer LLMs. They are causally extractable from a small set of attention heads and remain robust across context shifts.

  • Discussion: Function vectors provide a common task-representation mechanism across tasks and contexts, despite the apparent complexity of LLM metalearning.The discussion describes them as vectors representing tasks within hidden states.
  • Discussion: A small fixed set of identifiable attention heads can be used to explicitly extract function vectors from transformer LLMs.The extracted vectors have strong causal effects when inserted into the model.

ETHICS

Greater transparency into function representation may make large language models easier to audit, understand, and control, but could also enable abuse.

  • The paper cautions that transparency into function mechanisms may enable bad actors to inject or amplify functions causing undesirable behavior.

A DISCUSSION: FUNCTION VECTORS VS SEMANTIC VECTOR ARITHMETIC

The paper argues that function vectors are distinct from semantic word-embedding offsets: they encode information beyond output vocabulary and trigger nonlinear function execution through middle-layer computations.

  • Causal evidence: Causal mediation analysis distinguishes function vectors from correlated task encodings by identifying attention heads whose activations transport the task representation.The paper uses activation patching to extract computationally active representations rather than relying only on correlations.
  • Cyclic mappings: Function vectors can implement cyclic mappings such as antonyms, which cannot be represented by a fixed semantic vector offset.The contradiction arises because an antonym offset would need to map each word to its antonym in both directions using the same nonzero vector.
  • Information beyond vocabulary: Function-vector performance cannot be recovered from the target output vocabulary alone, although matching enough vocabulary words can recover part of it.Reconstructions with the same decoder vocabulary distribution perform poorly when only the top 100 words are matched, while broader matches sometimes recover a fraction of performance.
  • Middle-layer mediation: Function-vector causal effects are strongest at early and middle layers but fall to near-zero at later layers across evaluated tasks and models.This layer pattern suggests that function vectors activate nonlinear middle-layer components that execute the task rather than fully executing it themselves.
  • Interpretation: The evidence supports treating function vectors as references to functions that directly trigger nonlinear execution of abstract functions.The vectors are therefore not simple word embeddings, embedding differences, vocabulary averages, or semantic probability adjustments.

C EXPERIMENTAL DETAILS

The experiments extract function vectors from task-conditioned attention-head activations and evaluate them by intervening on hidden states across prompts, layers, and tasks.

  • Function-vector extraction: A function vector is computed from task-conditioned mean activations of selected attention heads, ranked by average indirect effect.For GPT-J, the extraction uses 100 clean 10-shot prompts, while head ranking uses 25 corrupted 10-shot prompts per task.
  • Function-vector extraction: GPT-J uses 10 attention heads because performance begins to plateau around |A| = 10 for a majority of 18 evaluated tasks.Some tasks continue to improve when more heads are included.
  • Evaluation: Function-vector causal effect is measured by adding the vector to a chosen layer’s last-token hidden state and scoring top-1 answer accuracy.If the expected answer has multiple tokens, evaluation targets its first token.
  • Evaluation: The default intervention occurs near one-third of the network depth: ℓ≈|L|/3, corresponding to layer 9 in GPT-J and model-specific layers in other models.The paper reports this location as working well in practice for most Section 3 and appendix results.
  • Prompt settings: Function vectors extracted from the default Q/A template are tested on shuffled-label, zero-shot, and other prompt templates.Zero-shot prompts contain only the query, without prepended examples, and Table 8 covers additional templates.

D RESULTS INCLUDING INCORRECT ICL

Removing correct-answer filtering leaves the causal effects of function vectors essentially unchanged. The evaluation spans diverse abstractive and extractive tasks, including generated, curated, and benchmark datasets.

  • Incorrect ICL Answers: Function-vector causal effects remain essentially unchanged when test cases are not filtered to correct 10-shot ICL answers.The comparison is reported in Figure 7.
  • Task Coverage: The evaluation covers diverse abstractive tasks, including question answering, classification, transformations, relational mappings, inflection, translation, and sentiment analysis.Examples include CommonsenseQA, AG News, capitalization, country-capital, present-past, translation, and SST-2 sentiment.
  • Dataset Construction: Several datasets require filtering or manual correction because source data contain duplicate answers, inaccurate generated entries, or incomplete examples.Antonym, synonym, translation, park-country, present-past, and sentiment datasets apply such processing.
  • Extractive Tasks: Extractive evaluation tests whether function vectors can support retrieval when the answer appears somewhere in the prompt.CoNLL-2003 subsets target person, location, and organization entities while restricting examples to sentences with a single specified entity.

G CAUSAL MEDIATION ANALYSIS

Causal mediation analysis identifies attention heads that mediate information flow during in-context learning. Across models, the most causally influential heads generally cluster in middle layers, with GPT-NeoX shifted earlier.

  • Cross-Model Causal Effects: For Llama 2 (7B), the highest-AIE heads have a lower maximum effect than GPT-J despite Llama 2 having more than twice as many attention heads.The reported maxima are approximately 0.047 for Llama 2 (7B) and 0.053 for GPT-J.
  • Cross-Model Causal Effects: GPT-NeoX is an exception, with its highest-AIE heads clustered in earlier middle layers, approximately layers 10–20.Other evaluated models place influential heads more toward the network middle.

H ATTENTION PATTERNS AND PREFIX-MATCHING SCORE

Causally influential heads attend most strongly to exemplar outputs, and some also match induction-head behavior. Function-vector effects are strongest in early-middle layers across model sizes, while decoded vocabularies only partly explain the vectors.

  • Attention Patterns: Across tasks, high-effect heads consistently place their strongest attention on the output tokens of each few-shot exemplar.This pattern appears across singular-plural, present-past, country-capital, and capitalization tasks.
  • Prefix Matching: Three of GPT-J’s ten highest-AIE heads also have high prefix-matching scores: 8-1, 12-10, and 24-6.Their prefix-matching scores are 0.49, 0.56, and 0.31, respectively.
  • Prefix Matching: Induction heads contribute to function-vector formation, but other causally influential heads do not show the same induction pattern on repeated random-token sequences.Thus, the identified causal set is broader than conventional induction heads.
  • FV Decoding: Function vectors often encode task output vocabularies, but output-vocabulary information alone does not fully identify a function vector.Sentiment shows strong output-space encoding, whereas translation does not encode output spaces as clearly.
  • Layer Placement: Across Llama 2 model sizes, adding function vectors generally produces the highest zero-shot accuracies in early-middle layers.This layer pattern persists regardless of the model’s total number of layers.

L EVALUATING FUNCTION VECTORS ON CYCLIC TASKS

Cyclic next-item and previous-item tasks test whether function vectors trigger relational behavior rather than act as simple semantic offsets. Function vectors substantially improve zero-shot performance, especially when inserted in early-middle layers.

  • Task Design: Next-item and previous-item tasks use reciprocal relations over cyclic and non-cyclic sequences, including weekdays, months, letters, numbers, and Roman numerals.Previous-item reverses the direction of the corresponding next-item pairs.
  • ICL Performance: With 10 examples, models perform both cyclic tasks fairly well, while next-item performance is usually higher than previous-item performance.Performance generally increases as more examples are provided.
  • Zero-Shot FV Effects: Function vectors significantly improve zero-shot accuracy over the baseline for both next-item and previous-item tasks across GPT-J and Llama 2 model sizes.Figure 29 compares dotted baseline accuracies with solid accuracies after FV addition.
  • Zero-Shot FV Effects: Peak cyclic-task performance occurs when the function vector is added in early-middle layers, followed by a sharp decline about two-thirds through the network.Example outputs show FVs handling boundary cases that the base model usually answers by copying the query.
  • Interpretation: Successful cyclic behavior provides additional evidence that function vectors are not merely simple semantic vector offsets.The argument extends from two-step antonym cycles to longer cycles such as the seven-day cycle.

M AN ALTERNATIVE EVALUATION OF FUNCTION VECTORS

The alternative intervention distributes task-conditioned head activations across their corresponding layers, whereas the original function vector concentrates them at one layer. Across zero-shot and shuffled-label settings, the alternative generally reproduces function-vector causal effects, while vocabulary-space analysis shows that function vectors promote query-specific plausible answers rather than generic output words.

  • Alternative intervention: The alternative approach adds each task-conditioned head activation to its corresponding layer, unlike the function vector, which adds all selected head outputs at one layer.The intervention is applied at every layer represented by the selected attention heads.
  • Causal comparison: In zero-shot evaluation, the alternative matches function-vector performance for most tasks, performs worse on English-French, and performs better on Present-Past.The comparison covers six representative tasks using GPT-J.
  • Causal comparison: In shuffled-label evaluation, the alternative matches the function vector’s causal effects for all tasks.The results compare both interventions against the base model across the six-task evaluation.
  • Causal comparison: The alternative typically achieves the same peak performance, but only the concentrated function-vector intervention reveals a performance dropoff around two-thirds through the network.This difference reflects the visibility of layer-specific effects under the two intervention designs.
  • Vocabulary-space effects: Function-vector intervention promotes tokens corresponding to plausible answers for each individual query rather than merely generic words from the task’s output space.Country-Capital, Antonym, and English-French examples show query-specific answer promotion, including valid alternatives even when top-1 accuracy is incorrect.
Loading 2310.15213v2…