MEMIT + Forward Pass Shift

This repo is a fork of MEMIT (Offical Repo) that includes the implementation of Forward Pass Shift from BackwardLens (Offical Repo)

Follow MEMIT instructions to run the experiments and see down below instruction regarding Forward Pass Shift

MEMIT: Mass-Editing Memory in a Transformer

Editing thousands of facts into a transformer memory at once.

Table of Contents

• Installation
• MEMIT Algorithm Demo
• Running the Full Evaluation Suite
• Generating Scaling Curves
• How to Cite

Installation

We recommend conda for managing Python, CUDA, and PyTorch; pip is for everything else. To get started, simply install conda and run:

CONDA_HOME=$CONDA_HOME ./scripts/setup_conda.sh

$CONDA_HOME should be the path to your conda installation, e.g., ~/miniconda3.

MEMIT Algorithm Demo

notebooks/memit.ipynb demonstrates MEMIT. The API is simple; simply specify a requested rewrite of the following form:

request = [
    {
        "prompt": "{} plays the sport of",
        "subject": "LeBron James",
        "target_new": {
            "str": "football"
        }
    },
    {
        "prompt": "{} plays the sport of",
        "subject": "Michael Jordan",
        "target_new": {
            "str": "baseball"
        }
    },
]

Other similar example(s) are included in the notebook.

Running the Full Evaluation Suite

experiments/evaluate.py can be used to evaluate any method in baselines/.

For example:

python3 -m experiments.evaluate \
    --alg_name=MEMIT \
    --model_name=EleutherAI/gpt-j-6B \
    --hparams_fname=EleutherAI_gpt-j-6B.json \
    --num_edits=10000 \
    --use_cache

Results from each run are stored at results/<method_name>/run_<run_id> in a specific format:

results/
|__ MEMIT/
    |__ run_<run_id>/
        |__ params.json
        |__ case_0.json
        |__ case_1.json
        |__ ...
        |__ case_10000.json

To summarize the results, you can use experiments/summarize.py:

python3 -m experiments.summarize --dir_name=MEMIT --runs=run_<run1>,run_<run2>

Running python3 -m experiments.evaluate -h or python3 -m experiments.summarize -h provides details about command-line flags.

Forward Pass Shift

To produce experiments that include Forward Pass Shift, run:

scripts_v2/run_basic_exp_eval.sh 1 cf 1000

How to Cite

First and foremost, we thank the authors of the MEMIT paper. Please cite them:

@article{meng2022memit,
  title={Mass Editing Memory in a Transformer},
  author={Kevin Meng and Sen Sharma, Arnab and Alex Andonian and Yonatan Belinkov and David Bau},
  journal={arXiv preprint arXiv:2210.07229},
  year={2022}
}

If you used Foward Pass Shift, please also cite BackwardLens paper with:

@article{katz2024backward,
  title={Backward lens: Projecting language model gradients into the vocabulary space},
  author={Katz, Shahar and Belinkov, Yonatan and Geva, Mor and Wolf, Lior},
  journal={arXiv preprint arXiv:2402.12865},
  year={2024}
}