Skip to content

Latest commit

 

History

History
44 lines (28 loc) · 2.95 KB

README.md

File metadata and controls

44 lines (28 loc) · 2.95 KB

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 :

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.

Illustration of the sampling process for DiffRS. The path with the green background represents the DiffRS sampling process, and the rightmost images are generated when the images are sampled as a base sampler without rejection from the intermediate image. Timesteps are expressed as the noise level σ from the EDM scheme.

Overview of DiffRS. We sequentially apply the rejection sampling on the pre-trained transition kernel (red) to align the true transition kernel (blue). The acceptance probability is estimated by the time-dependent discriminator.

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