Official implementation of Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors (NeurIPS 2023).

Please feel free to email Tom or raise an issue with this repository and we'll get back to you as soon as possible.

Setup

1. Create a virtual environment (we use conda)

   conda env create --name grace_env --file environment.yml
   

2. Activate the virtual environment

   conda activate grace_env
   

3. Install the repository

   pip install -e .
   

Data

The QA experiments use data linked by the MEND repository. Per their instructions, you can download the data for NQ and zsRE from their Google Drive link and unzip each sub-directory into grace/data. SCOTUS and Hallucination data are handled through huggingface. zip files can be unzipped using tar -xf [filename].zip.

The data for the SCOTUS experiments is available on huggingface

Models

Our pre-trained models are public:

• BERT finetuned for SCOTUS
• GPT2 finetuned for Hallucination

Running experiments

Experiments are run using main.py. Experiment settings and hyperparameters are chosen using hydra. While more examples are available in ./scripts/main.sh, three representative experiments can be run as follows:

Editing GPT2-XL on Hallucination with GRACE

python grace/main.py experiment=hallucination model=gpt2xl editor=grace

Editing BERT on SCOTUS with GRACE

python grace/main.py experiment=scotus model=bert editor=grace

Editing T5 on zsRE with GRACE

python grace/main.py experiment=qa model=t5small editor=grace

Repository Roadmap

• ./scripts/ contains handy shell scripts for starting and running experiments in slurm.
• ./notebooks/ contains a simple example of editing a model with GRACE.
• ./ckpts/ will contain checkpoints of your edited models if you choose to checkpoint models.
• ./data/ will contain downloaded datasets if you choose to cache data yourself instead of relying on HuggingFace.
• ./grace/ contains the source code to GRACE
  • ./grace/main.py is the main file to kick off experiments.
  • ./grace/config/ contains the config files for datasets, editors, and pretrained models.
  • ./grace/editors/ contains source code for each compared editor.
  • ./grace/dataset.py contains source code for each compared dataset.
  • ./grace/metrics.py contains source code for each compared dataset.
  • ./grace/models.py contains source code for loading pretrained models.

Citation

Please use the following to cite this work:

@inproceedings{hartvigsen2023aging,
  title={Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors},
  author={Hartvigsen, Thomas and Sankaranarayanan, Swami and Palangi, Hamid and Kim, Yoon and Ghassemi, Marzyeh},
  booktitle={Advances in Neural Information Processing Systems},
  year={2023}
}