New sampler: Restart

[drax-xard]

Restart sampling

Generative processes that involve solving differential equations, such as diffusion models, frequently necessitate balancing speed and quality. ODE-based samplers are fast but plateau in performance while SDE-based samplers deliver higher sample quality at the cost of increased sampling time. We attribute this difference to sampling errors: ODE-samplers involve smaller discretization errors while stochasticity in SDE contracts accumulated errors. Based on these findings, we propose a novel sampling algorithm called Restart in order to better balance discretization errors and contraction.

Empirically, Restart sampler surpasses previous diffusion SDE and ODE samplers in both speed and accuracy. Restart not only outperforms the previous best SDE results, but also accelerates the sampling speed by 10-fold / 2-fold on CIFAR-10 / ImageNet
. In addition, it attains significantly better sample quality than ODE samplers within comparable sampling times. Moreover, Restart better balances text-image alignment/visual quality versus diversity than previous samplers in the large-scale text-to-image Stable Diffusion model pre-trained on LAION .

Paper: https://arxiv.org/abs/2306.14878
Github: https://github.com/newbeeer/diffusion_restart_sampling
Sampler implementation code: https://github.com/Newbeeer/diffusion_restart_sampling/blob/b1e0022406e619cafb118a69b70a90484fb84c94/diffuser/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L676-L817