This repository contains simplified code for the paper:
Corrective Machine Unlearning
Shashwat Goel*, Ameya Prabhu*, Philip Torr, Ponnurangam Kumaraguru, Amartya Sanyal
* = Equal Contribution
- 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 themanip_dataset
function andDatasetWrapper
class. - To add new unlearning methods, modify
src/methods.py
- To add new datasets, modify the
load_dataset
function insrc/datasets.py
- In general, any of these changes would require updates to
src/main.py
,src/opts.py
andsrc/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.
We hope you are excited by the potential for corrective unlearning! To cite our work:
@article{goel2024corrective,
title={Corrective Machine Unlearning},
author={Shashwat Goel and Ameya Prabhu and Philip Torr and Ponnurangam Kumaraguru and Amartya Sanyal},
journal={Transactions on Machine Learning Research},
year={2024},
url={https://openreview.net/forum?id=v8enu4jP9B}
}