Generalised Parallel Tempering: Flexible Replica Exchange via Flows and Diffusions

Leo Zhang, Peter Potaptchik, George Deligiannidis, Arnaud Doucet, Hai-Dang Dau, Saifuddin Syed

Keywords: MCMC, Normalising Flows, Diffusion Models

TL;DR: We introduce and formalise the design of more flexible swap kernels within the classical MCMC algorithm of Parallel Tempering to allow for the use of normalising flows and diffusion models to accelerate sampling.

Abstract: Parallel Tempering is a classical MCMC algorithm designed for leveraging parallel computation to sample efficiently from high-dimensional, multimodal or otherwise complex distributions via annealing. One limitation of the standard formulation of Parallel Tempering is the growth of computational resources required to generate high-quality samples, as measured by effective sample size or round trip rate, for increasingly challenging distributions. To address this issue, we propose the framework: Generalised Parallel Tempering (GePT) which allows for the incorporation of recent advances in modern generative modelling, such as normalising flows and diffusion models, within Parallel Tempering, while maintaining the same theoretical guarantees as MCMC-based methods. For instance, we show that this allows us to utilise diffusion models in a parallelised manner, bypassing the usual computational cost of a large number of steps to generate quality samples. Further, we empirically demonstrate that GePT can improve sample quality and reduce the growth of computational resources required to handle complex distributions over the classical algorithm.

Submission Number: 63