memit

Introduction to MEMIT: Mass-Editing Memory in a Transformer

MEMIT is an innovative project focused on editing large volumes of information, or "facts," in the memory of a transformer model simultaneously. The project offers a comprehensive framework for altering these facts, making it possible to dynamically update and control what a model remembers.

Installation

To get started with MEMIT, it is recommended to use conda, a package manager, to handle Python, CUDA, and PyTorch dependencies. Beyond those core dependencies, pip is used to install any additional packages. By setting up your environment using conda, as shown in the script below, users can ensure that they have all necessary tools:

CONDA_HOME=$CONDA_HOME ./scripts/setup_conda.sh

In the command above, $CONDA_HOME should be the directory path where your conda is installed—such as ~/miniconda3.

MEMIT Algorithm Demo

The MEMIT project provides an interactive demonstration via a Jupyter Notebook found in notebooks/memit.ipynb. The demonstration showcases how to use the MEMIT API to execute what is called a "requested rewrite". The example below shows how users can modify specific facts:

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"
        }
    },
]

In this example, the memory associated with the subjects (LeBron James and Michael Jordan) is updated with new target information (football and baseball, respectively).

Running the Full Evaluation Suite

To assess the performance of the MEMIT algorithm or any alternative methods available in the baselines/ directory, users can utilize the experiments/evaluate.py script. An example command for running an evaluation is:

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

The execution results are organized into directories by method name and run ID, structured for detailed analysis. As seen in the example below, each result is stored in its own JSON file within these directories:

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

To further analyze these results, the experiments/summarize.py script can be employed:

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

Both the evaluation and summarization scripts offer command-line help functionality for more detailed instructions on their usage.

How to Cite

If you’re referencing the MEMIT work in academic contexts, it’s been documented under the following citation:

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

With MEMIT, the possibility of rapidly updating the memory of transformer models opens new doors for research and application development, making it an exciting tool for anyone involved in machine learning and artificial intelligence.