The implementation of the paper "Efficient Fine-Tuning and Concept Suppression for Pruned Diffusion Models"

Comparison of generative quality and style adherence: Row 1: The original Stable Diffusion 2.1 model. Row 2: A pruned version fine-tuned with 20,000 iterations of combined DDPM and distillation loss. Row 3: A pruned version fine-tuned with 20,000 iterations of our proposed bilevel fine-tuning approach, removing styles of Van Gogh, Monet, and Picasso. Our bilevel method is successful in retaining generative quality and style diversity while suppressing undesirable concepts.

Why can a two-stage approach (fine-tuning followed by forgetting) be suboptimal? If fine-tuning yields \(\hat{\theta}\), initializing the concept unlearning parameters with \(\hat{\theta}\) and optimizing the concept unlearning loss results in \(\theta'\), which is suboptimal for both fine-tuning and for concept unlearning. In contrast, our bilevel method produces the optimal solution \(\theta^*\), achieving better performance for both fine-tuning and unlearning.

Installation

Follow these steps to set up the project:

1. Create Conda Environment

Use the provided env.yaml file:

conda env create -f env.yaml

2. Activate the Conda Environment

Activate the environment:

conda activate unlearn-ft

3. Install Project Dependencies

From the project root directory, install the dependencies:

pip install -e .

Data Preparation

Prepare the data for training as mentioned in the paper. You can also adapt our method for your own dataset with minor code modifications.

1. Download MS-COCO 2017

1.1 Download the training and validation images

Download 2017 train and 2017 val images from the COCO website. Place them in your chosen directory.

1.2 Download the annotations

Download the 2017 train/val annotations and place them in the same directory as the images. Your directory should look like this:

coco
├── annotations
│   ├── captions_train2017.json
│   ├── captions_val2017.json
│   └── ...
└── images
    ├── train2017
    │   ├── COCO_train2017_000000000009.jpg
    │   └── ...
    └── val2017
        ├── COCO_val2017_000000000042.jpg
        └── ...

Training

1. Effect of Pruning

There are pruning checkpoints (55 and 82) corresponding to the paper experiments in this HF Repo. You can download them and put them in a path of your choice. Then, to reproduce the experiments on the effect of pruning, you can use the following slurm scripts. Make sure you update the checkpoint paths in the scripts. In all scripts, set expert ID = 2 for the 55% checkpoint and expert ID = 5 for the 82% checkpoint if expert id is required.

slurm_scripts/coco/baseline_coco_aptp_ddpm_512.slurm
slurm_scripts/coco/baseline_coco_aptp_ddpm_512_random.slurm

and similarly for other loss variants (distillation, both) if you're interested.

2. Effect of Distillation

To reproduce the experiments on the effect of distillation, you can use the following slurm scripts:

slurm_scripts/coco/baseline_coco_aptp_ddpm_512.slurm
slurm_scripts/coco/baseline_coco_aptp_distillation_512.slurm
slurm_scripts/coco/baseline_coco_aptp_both_512.slurm

These scripts produce the results for 55. You can modify them to reproduce the results for 82 as well. Just change the config and ckpt paths.

3. Bilevel Finetuning

To reproduce the bilevel finetuning experiments, you can use the following slurm scripts:

slurm_scripts/coco/baseline_coco_aptp_both_512_bilevel.slurm (For Monet. You can simply change the config file for other desired concepts).
slurm_scripts/coco/baseline_coco_aptp_both_512_bilevel_nudity.slurm

4. Baselines

The code for reproducing the baseline results is provided in the baselines folder. See the run.sh file in each baseline folder.

Image Generation

To generate images from any finetuned checkpoint, use the slurm scripts:

slurm_scripts/evaluation/gen_images.slurm path/to/ft/ckpt

Evaluation

To evaluate Unlearn-FT, we report the FID, CLIP Score.

1. FID Score

We use clean-fid to calculate the FID score. The numbers reported in the paper are calculated using this pytorch legacy mode.

1.1 MS-COCO Preparation

We sample use the validation images from the 2017 MS-COCO. Check out the sample and resize script.

1.2 Generate Custom Statistics

Generate custom statistics for both sets of reference images::

from cleanfid import fid
fid.make_custom_stats(dataset, dataset_path, mode="legacy_pytorch") # mode can be clean too.

Now we can calculate the FID score for the generate images using the provided script.

2. CLIP Score

To calculate clip score, we use this library. Extract features of reference images with the clip feature extraction script and calculate the score using the clip score script.

3. Unlearning Benchmarks

To evaluate the models on unlearing benchmarks, follow their instructions and use the checkpoints or generated images from the experiments here. See CSD, Artist Erasure Eval (CP), and NSFW removal Eval.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you find this work useful, please consider citing the following paper:

@inproceedings{shirkavand2025efficient,
  title={Efficient fine-tuning and concept suppression for pruned diffusion models},
  author={Shirkavand, Reza and Yu, Peiran and Gao, Shangqian and Somepalli, Gowthami and Goldstein, Tom and Huang, Heng},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={18619--18629},
  year={2025}
}

Table of Contents