Issues:
We propose
In this project, we propose a novel Separable, Recoverable, and Sustainable Multi-concept Eraser (SRS-ME), enabling diffusion models to forget all concepts that they should forget without necessitating retraining from scratch. Specifically, through theoretical analysis, we introduce the paradigm of weight decoupling for constructing separable weight shifts, which can decouple interactions among weight shifts targeting diverse concepts. This approach also provides flexibility in both erasing and recovering arbitrary concepts while preserving model watermarks. To effectively erase inappropriate concepts and preserve model performance on regular concepts, we design an innovative concept-irrelevant unlearning optimization process. By defining concept representations, this process introduces the concept correlation loss and the momentum statistic-based stopping condition. Besides, to reduce memory usage, we demonstrate the feasibility of optimization decoupling for separated weight shifts. Benchmarked against prior work, extensive experiments demonstrate
We do not offer the fine-tuned weights for download, as the training process is faster.
stable-diffusion
main directory of stable diffusion. stable-diffusion/models/ldm/
After installation, follow these instructions to train a machine unlearning model:
./data/train/{0,1,2,3,4,5,6,7,8,9}
or ./data/eval/{0,1,2,3,4,5,6,7,8,9}
python eval-scripts/generate-data.py --prompt 'Van Gogh' --model_name '0' --save_path './data/' --num_samples 1 --device 'cuda:0'
--prompt
should be Van Gogh
, Picasso
, Cezanne
, Jackson Pollock
, Caravaggio
, Keith Haring
, Kelly McKernan
, Tyler Edlin
, and Kilian Eng
for style unlearning.
Train classification model
python train-scripts/artist-cls-train.py --device 'cuda:0'
Train comparative experiments: FMN, Abconcept, Esd
python train-scripts/FMN.py --erase_cat 'style' --erased_index 0 --lr 1e-5 --iter_break 50 --devices '0,1'
python train-scripts/Abconcept.py --erase_cat 'style' --erased_index 0 --lr 1e-5 --iter_break 75 --devices '0,1'
python train-scripts/Esd.py --iter_break 1000 --train_method 'xattn' --erase_cat 'style' --erased_index 0 --devices '0,1'
Train our SRS-ME
Calculate particular solutions:
python train-scripts/SRSME_Solution.py --erase_cat 'style' --erased_index 0 --em_indexes '1,2' --scenei 1 --timei 0 --device '0'
SRS-ME optimization:
python train-scripts/SRS_ME.py --lr 1e-1 --scale_factor 1e-4 --threshold 1e-4 --reg_beta 1e-6 --erased_index 0 --scenei 1 --timei 0 --max_base 691 --erase_cat 'style' --devices '0,1'
To generate images from one of the custom models use the following instructions:
eval-scripts/generate-images-merge.py
you would need a csv file with columns prompt
, evaluation_seed
and case_number
. (Sample data in data/
)num_samples
) using trained DMs per seed. For instance, select 0, 1, 2 model weights in generate-images-merge.py, and yield 012-merge
python eval-scripts/generate-images-merge.py --prompts_path './data/style_classify.csv' --save_path 'evaluation_folder' --num_samples 5 --device 'cuda:0'
To obtain the quantitative results:
Rename the dir_name of generated images to eval
, cate-cls-test.py
for object, artist-cls-test.py
for artist style
python train-scripts/move_image.py
python train-scripts/cate-cls-test.py
The default
contains the images generated by the original diffusion model; The eval
contains the images generated by the unlearned diffusion model
python eval-scripts/lpips_eval.py --path1 './evaluation_folder/ori-merge' --path2 './evaluation_folder/012-merg'
If you find our work valuable, please consider citing the preprint using the following format: