maciek-pioro / function_vectors

Function Vectors in Large Language Models (ICLR 2024)

Function Vectors in Large Language Models

Project Website | Arxiv Preprint

This repository contains data and code for the paper: Function Vectors in Large Language Models.

Setup

We recommend using conda as a package manager. The environment used for this project can be found in the fv_environment.yml file. To install, you can run:

conda env create -f fv_environment.yml
conda activate fv

Demo Notebook

Checkout notebooks/fv_demo.ipynb for a jupyter notebook with a demo of how to create a function vector and use it in different contexts.

Data

The datasets used in our project can be found in the dataset_files folder.

Code

Our main evaluation scripts are contained in the src directory with sample script wrappers in src/eval_scripts.

Other main code is split into various util files:

• eval_utils.py contains code for evaluating function vectors in a variety of contexts
• extract_utils.py contains functions for extracting function vectors and other relevant model activations.
• intervention_utils.py contains main functionality for intervening with function vectors during inference
• model_utils.py contains helpful functions for loading models & tokenizers from huggingface
• prompt_utils.py contains data loading and prompt creation functionality

Citing our work

This work appeared at ICLR 2024. The paper can be cited as follows:

@inproceedings{todd2024function,
    title={Function Vectors in Large Language Models}, 
    author={Eric Todd and Millicent L. Li and Arnab Sen Sharma and Aaron Mueller and Byron C. Wallace and David Bau},
    booktitle={Proceedings of the 2024 International Conference on Learning Representations},
    year={2024},
}