drimpossible / corrective-unlearning-bench

Corrective Unlearning

This repository contains simplified code for the paper:

Corrective Machine Unlearning
Shashwat Goel*, Ameya Prabhu*, Philip Torr, Ponnurangam Kumaraguru, Amartya Sanyal

* = Equal Contribution

[PDF] [Slides] [Bibtex]

Usage

• The code first pretrains a model if one with the specified parameters is not alredy available. Otherwise it uses the stored model. It then performs the specified unlearning procedure.

$ python main.py --dataset=CIFAR10 --num_classes=10 --model=resnet9 --pretrain_iters=4000 --dataset_method=labeltargeted --forget_set_size=500 --deletion_size=250 --unlearn_method=EU --unlearn_iters=4000 --k=-1

The above script trains a Resnet9 model on CIFAR10 with 500 samples manipulated using Interclass Confusion. It then picks 250 samples known to the unlearning procedure, here retraining the whole model from scratch without the deletion data.
Arguments you can freely tweak given a dataset and model:

• Pretrain Iterations (--pretrain_iters)
• Manipulation Type (--dataset_method)
• No. of Manipulated Samples (--forget_set_size)
• No. of unlearning samples known to developers (--deletion_size)
• Unlearning method (--unlearn_method)
• Method specific hyperparameters, which can be inferred from src/opts.py

Adding functionality can be done as follows:

• To add new evaluations, modify src/datasets.py, specifically the manip_dataset function and DatasetWrapper class.
• To add new unlearning methods, modify src/methods.py
• To add new datasets, modify the load_dataset function in src/datasets.py
• In general, any of these changes would require updates to src/main.py, src/opts.py and src/visualize.py

Additional details and default hyperparameters can be found in src/opts.py

• To replicate the complete set of experiments, run scripts/CIFAR10_poisoning.sh and similar for other datasets and manipulations.

If you discover any bugs in the code please contact me, I will cross-check them with my nightmares.

Citation

We hope you are excited by the potential for corrective unlearning! To cite our work:

@misc{goel2024corrective,
      title={Corrective Machine Unlearning}, 
      author={Shashwat Goel and Ameya Prabhu and Philip Torr and Ponnurangam Kumaraguru and Amartya Sanyal},
      year={2024},
      eprint={2402.14015},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}