Dlut-lab-zmn / SRS-ME

Separable Diffusion Model Unlearning
12 stars 0 forks source link

Forget All That Should Be Forgotten: Separable, Recoverable, and Sustainable Multi-Concept Erasure from Diffusion Models

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

Fine-tuned Weights

We do not offer the fine-tuned weights for download, as the training process is faster.

Installation Guide

Training Guide

After installation, follow these instructions to train a machine unlearning model:

  1. Generate data and then put these samples to ./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}

--prompt should be Van Gogh, Picasso, Cezanne, Jackson Pollock, Caravaggio, Keith Haring, Kelly McKernan, Tyler Edlin, and Kilian Eng for style unlearning.

  1. Train classification model

    • python train-scripts/artist-cls-train.py --device 'cuda:0'
  2. 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'
  3. Train our SRS-ME

Calculate particular solutions:

SRS-ME optimization:

Generating Images

To generate images from one of the custom models use the following instructions:

For instance, select 0, 1, 2 model weights in generate-images-merge.py, and yield 012-merge

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

The default contains the images generated by the original diffusion model; The eval contains the images generated by the unlearned diffusion model