# Dual-Model Defense against Membership Inference Attacks on Diffusion Models

This repository contains the code for the paper ["Dual-Model Defense: Safeguarding Diffusion Models from Membership Inference Attacks through Disjoint Data Splitting"]().

We acknowledge the use of the code from [SecMIA](https://github.com/jinhaoduan/SecMI), [PIA](https://github.com/kong13661/PIA), [DCR](https://github.com/somepago/DCR), and [Diffusion Memorization](https://github.com/YuxinWenRick/diffusion_memorization).

The Pokemon dataset can be downloaded from [here](https://drexel0-my.sharepoint.com/:u:/g/personal/jd3734_drexel_edu/EeEwxOQ-5cZEnf534S6WRkQBOcvbAtfmuV-h5UjyIF8YxQ?e=JYAHfo). Thanks [SecMIA](https://github.com/jinhaoduan/SecMI-LDM) for providing the dataset. All the downloaded data should be put in the `./datasets` folder

## How to run

### Environment Setup
``` bash
conda create -n ddmd python=3.10
conda activate ddmd
pip install -r requirements.txt
```

### Model Training
DDPM:
``` bash
cd DDPM
# Train a normal DDPM model
bash scripts/train.sh
# Train two models with disjoint data
bash scripts/train_dual.sh
# Distill to a private student model
bash scripts/distill.sh
```

LDM:
``` bash
cd LDM
# Generate captions for the dataset
bash scripts/generate-captions.sh
# Train LDM with prompt diversification
bash scripts/train-ldm.sh
# Train two models with disjoint data
bash scripts/train-dual.sh
# Distill to a private student model
bash scripts/distill-ldm.sh
```

### White-box Attack
DDPM:
``` bash
cd DDPM
# SecMIA
bash scripts/secmia.sh
# PIA
bash scripts/pia.sh
```

LDM:
``` bash
cd LDM
# PIA
bash scripts/pia.sh
```
The SecMI attack for LDM is conducted following [SecMIA](https://github.com/jinhaoduan/SecMI-LDM).


### Black-box Attack
``` bash
cd black-box-attack
bash attack.sh
```

### Memorization Detection
``` bash
pip install diffusers==0.18.2
cd memorization
# Extract scores with with t-err loss
bash detect_mem.sh
```
For detection visualization, follow `./memorization/examples/det_mem_viz.ipynb`.

## References

Duan, J., Kong, F., Wang, S., Shi, X., & Xu, K. (2023, July). Are diffusion models vulnerable to membership inference attacks?. In International Conference on Machine Learning (pp. 8717-8730). PMLR.

Kong, F., Duan, J., Ma, R., Shen, H. T., Shi, X., Zhu, X., & Xu, K. An Efficient Membership Inference Attack for the Diffusion Model by Proximal Initialization. In The Twelfth International Conference on Learning Representations.

Somepalli, G., Singla, V., Goldblum, M., Geiping, J., & Goldstein, T. (2023). Understanding and mitigating copying in diffusion models. Advances in Neural Information Processing Systems, 36, 47783-47803.