Journey to the BAOAB-limit: finding effective MCMC samplers for score-based models

Ajay Jain, Ben Poole

Keywords: Diffusion, MCMC, Sampling

TL;DR: Sometimes bugs are effective MCMC samplers for score-based models.

Abstract: Diffusion and score-based generative models have achieved remarkable sample quality on difficult image synthesis tasks. Many works have proposed samplers for pretrained diffusion models, including ancestral samplers, SDE and ODE integrators and annealed MCMC approaches. So far, the best sample quality has been achieved with samplers that use time-conditional score functions and move between several noise levels. However, estimating an accurate score function at many noise levels can be challenging and requires an architecture that is more expressive than would be needed for a single noise level. In this work, we explore MCMC sampling algorithms that operate at a single noise level, yet synthesize images with acceptable sample quality. We show that while naïve application of Langevin dynamics and a related noise-denoise sampler produces poor samples, methods built on integrators of underdamped Langevin dynamics using splitting methods can perform well. Our samplers also have great diversity, allowing many samples to be generated in a single long-run MCMC chain. Further, by combining MCMC methods with existing multiscale samplers, we begin to approach competitive sample quality without using scores at large noise levels. Find videos and code at https://ajayj.com/journey.

Student Paper: Yes