Diffusion Variational Inference: Diffusion Models as Expressive Variational Posteriors

We propose denoising diffusion variational inference (DDVI), an approximate inference algorithm for latent variable models which relies on diffusion models as expressive variational posteriors. Our method augments variational posteriors with auxiliary latents, which yields an expressive class of models that perform diffusion in latent space by reversing a user-specified noising process. We fit these models by optimizing a novel lower bound on the marginal likelihood inspired by the wake-sleep algorithm. Our method is easy to implement (it fits a regularized extension of the ELBO), is compatible with black-box variational inference, and outperforms alternative classes of approximate posteriors based on normalizing flows or adversarial networks. When applied to deep latent variable models, our method yields the denoising diffusion VAE (DD-VAE) algorithm. We use this algorithm on a motivating task in biology--inferring latent ancestry from human genomes--outperforming strong baselines on the Thousand Genomes dataset. Latent variable methods often rely on variational inference to fit an approximate model of the posterior distribution (Vahdat & Kautz, 2020; Maaløe et al., 2016). The expressivity of this model has a significant impact on the performance of variational inference (Kingma et al., 2016), which motivates research that leverages modern generative models--including normalizing flows (Rezende & Mohamed, 2015) and generative adversarial networks (Goodfellow et al., 2014; Makhzani et al., 2015)--to represent expressive approximate posteriors.