| 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.

The requirements for this code are the same as DG.

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

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/discriminator/cifar_uncond/discriminator_60.pt' at DG.
• Download 'DG/checkpoints/ADM_classifier/32x32_classifier.pt' at DG.

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

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.

@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},
}