aailabkaist / DiffRS

Official PyTorch implementation for Diffusion Rejection Sampling (DiffRS) in ICML 2024.

Diffusion Rejection Sampling (DiffRS) (ICML 2024)

| paper | arXiv | poster |

---

This repo contains an official PyTorch implementation for the paper "Diffusion Rejection Sampling" in ICML 2024.

Byeonghu Na, Yeongmin Kim, Minsang Park, Donghyeok Shin, Wanmo Kang, and Il-Chul Moon

---

This paper introduces Diffusion Rejection Sampling (DiffRS), a new diffusion sampling approach that ensures alignment between the reverse transition and the true transition at each timestep.

Requirements

The requirements for this code are the same as DG.

In our experiment, we utilized CUDA 11.4 and PyTorch 1.12.

Diffusion Rejection Sampling

1. Download the pre-trained diffusion network and the trained discriminator network from DG.

• Download 'edm-cifar10-32x32-uncond-vp.pkl' at EDM.
• Download 'DG/checkpoints/ADM_classifier/32x32_classifier.pt' at DG.
• Download 32x32_classifier.pt at ADM.

2. Generate DiffRS samples using generate_diffrs.py. For example:

python3 generate_diffrs.py \
    --network checkpoints/pretrained_score/edm-cifar10-32x32-uncond-vp.pkl \
    --outdir=samples/cifar10/diffrs --rej_percentile=0.75 --max_iter=105

Acknowledgements

This work is heavily built upon the code from:

• Kim, D., Kim, Y., Kwon, S. J., Kang, W., & Moon, I. C. (2023, July). Refining Generative Process with Discriminator Guidance in Score-based Diffusion Models. In International Conference on Machine Learning (pp. 16567-16598). PMLR.

Citation

@inproceedings{na2024diffusion,
  title = 	 {Diffusion Rejection Sampling},
  author =       {Na, Byeonghu and Kim, Yeongmin and Park, Minsang and Shin, Donghyeok and Kang, Wanmo and Moon, Il-Chul},
  booktitle = 	 {Proceedings of the 41st International Conference on Machine Learning},
  pages = 	 {37097--37121},
  year = 	 {2024},
  editor = 	 {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
  volume = 	 {235},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {21--27 Jul},
  publisher =    {PMLR},
}