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//run_ in a specific format:

results/
|__ MEMIT/
    |__ run_/
        |__ 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_,run_

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

How to Cite

@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}
}