Code and documentation for the paper Not All Data Are Unlearned Equally

conda create -n unlearn python=3.10
conda activate unlearn
pip install -r requirements.txt

Make sure you set up the wandb configurations in the configuration files for training and unlearning.

To run the probabilistic experiments for OLMo (MMLU, Hellaswag), make sure that lm-eval-harness is installed in this directory. The evaluation script calls it as

python -m lm_eval --options

You can also turn these evaluations off by editing config/unlearn.yaml

1. We release the code to generate biography datasets for the GPT-2 experiments. The datasets can be created by running the biographies.ipynb notebook, which randomly samples new biographies and uploads them to huggingface. We refrain from publically releasing the biographies from the paper following the [privacy concerns] raised by the original authors.
2. The OLMo datasets, seperated by count buckets are released here.

the gpt2_bios.py and the config/gpt_train.yaml file contains the code for finetuning GPT2 on the fake biographies and questions. For GPU-parallel training, use

torchrun --nproc_per_node=4 gpt2_bios.py gpt2=124M LEARNING_RATE=1e-3 SCALE_FACTOR=10 EPOCHS=100 BATCH_SIZE=32 RANDOM_SEED=1234

Make sure that the model identifier is correctly provided in the config/model_config.yaml file. See config/datasets paths to add new datasets or use the datasets from the paper. You should specify a forget and retain split in the hyperparameters. The scripts expct that each split has a 'question' and 'answer' column with the unlearn/retain entries. To unlearn a model on a forget set, use the following command:

# python unlearn_train_eval.py \
#   --mask_retain_question False \
#   --loss_type SIMNPO_GD \
#   --dataset books \
#   --dataset.forget_split forget_high_count \
#   --dataset.retain_split forget_medium_count \
#   --model_family olmo-7b-sft \
#   --learning_rate 5e-6 \
#   --batch_size 8 \
#   --forget_loss_lambda 1.0 \
#   --retain_loss_lambda 5.0 \
#   --simnpo_gamma 0 \
#   --npo_beta 0.1 \
#   --unlearn_epochs 5 \
#   --seed 1`

You can choose between SIMNPO, GA and IDK for the loss function. Results are automatically uploaded to wandb.

The unlearn_train.py script automatically runs evaluation after each epoch. Incase you wish to make things faster by evaluating things in parallel, unlearn_eval.py accepts a directory and evaluates the model at this directory. The path should contain a yaml file with all necessary information to load the model, the dataset and the evaluation tasks. Example command:

python unlearn_eval.py --checkpoint_dir path/to/checkpoint --checkpoint_step 100