Diffusion Rejection Sampling (DiffRS) (ICML 2024)

| paper | arXiv | poster |

This repository is the implementation of Diffusion Rejection Sampling, for Gaudi-v2 of the code below :

• Na, B., Kim, Y., Park, M., Shin, D., Kang, W., & Moon, I. C. Diffusion Rejection Sampling. In Forty-first International Conference on Machine Learning.

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 experiments, we utilized the Docker image vault.habana.ai/gaudi-docker/1.17.0/ubuntu22.04/habanalabs/pytorch-installer-2.3.1:latest.

To run the Docker container based on this image, execute the following command:

docker pull vault.habana.ai/gaudi-docker/1.17.0/ubuntu22.04/habanalabs/pytorch-installer-2.3.1:latest

This command starts a new container from the specified image, providing an environment with the necessary dependencies for your project.

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