Rebuilding ROME : Resolving Model Collapse during Model Editing

Changes to the update equation

We focus on the way the MLP keys ($`k`$ such that $`Wk=v`$) are computed. See the rome/compute_u.py and rome/compute_v.py scripts for details. The derived ROME update equation is

$\hat{W}=W+\Lambda_* \left(C^{-1} k_*\right)^T$

where

$\Lambda_* =\frac{v_* - Wk_* }{ (C^{-1} k_* )^T k_* }$

$k_* =\frac{1}{N} \sum_{j=1}^N k\left(x_j+s\right).$

k(x)= \sigma \left( W_{f c}^{\left( l^* \right)} \gamma\left( a_{\[x\], i}^{\left( l^* \right)} + h_{\[x\], i}^{\left( l^* -1 \right)} \right)\right)

Note that the optimization step to compute $`v_*`$ is based on $`k_*`$. The original ROME implementation, however, computes

$\hat{W}=W+\Lambda\left(C^{-1} k_*\right)^T$

where

$\Lambda=\frac{v_* -Wk}{(C^{-1} k)^T k}$ $k = k(s)$

We find that the latter leads to rapid degradation in model performance in a sequential editing setting, and prone to particular edits known as disabling edits that render the model unusable post-update. Our experiments focus on unifying the computation of the keys in the update equation, and we study the use of $`k`$ and $`k_*`$.

Installation

We recommend using Docker to set up a clean dev environment.

docker compose up -d --build

To download the datasets used for evaluation, install Git LFS if needed:

curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install -y git-lfs
git lfs pull

Running the experiments

The script supports sequential editing with the --sequential flag. With sequential editing, the edited model is evaluated for downstream task performance on 4 GLUE datasets after every 20 edits. The interval can be changed within the code-base.

You can evaluate either GPT2-XL or GPTJ-6B using the appropriate hyperparameter file to configure how the update equation is computed.

python experiments/evaluate.py \
    --model_name=${MODEL_NAME} \
    --hparams_fname=${HPARAM_FILE_NAME} \
    --ds_name=cf \
    --sequential

How to Cite

If you find our work useful, please cite it using the following:

@article{gupta2024rebuilding,
  title={Rebuilding ROME: Resolving Model Collapse during Sequential Model Editing},
  author={Gupta, Akshat and Anumanchipalli, Gopala},
  journal={arXiv preprint arXiv:2403.07175},
  year={2024}
}

@article{gupta2024model,
  title={Model Editing at Scale leads to Gradual and Catastrophic Forgetting},
  author={Gupta, Akshat and Rao, Anurag and Anumanchipalli, Gopala},
  journal={arXiv preprint arXiv:2401.07453},
  year={2024}
}