Diffusion Priors In Variational Autoencoders
Abstract
Among likelihood-based approaches for deep generative modelling, variational autoencoders (VAEs) offer scalable amortized posterior inference and fast sampling.
However, VAEs are also more and more outperformed by competing models such as normalizing ﬂows (NFs), deep-energy models, or the new denoising diffusion probabilistic models (DDPMs).
In this preliminary work, we improve VAEs by demonstrating how DDPMs can be used for modelling the prior distribution of the latent variables.
The diffusion prior model improves upon Gaussian priors of classical VAEs and is competitive with NF-based priors.
Finally, we hypothesize that hierarchical VAEs could similarly beneﬁt from the enhanced capacity of diffusion priors.
1. Introduction
Over the last few years, the interest of the deep learning community for generative modelling has increased steadily.
Among the likelihood-based approaches for deep generative modelling variational autoencoders (Kingma & Welling, 2013, VAEs) stand as one of the most popular, although competing approaches now demonstrate better performance.
In particular, Ho et al.
(2020); Nichol & Dhariwal (2021);
Dhariwal & Nichol (2021) recently showed that denoising diffusion probabilistic models (DDPMs) are competitive deep generative models, obtaining samples quality similar to those of the best implicit deep generative models such as ProgressiveGAN (Karras et al., 2017) and StyleGAN (Karras et al., 2019).
Similarly to VAEs, DDPMs train on a variational bound and may be interpreted under the encoding-decoding framework.
In the original formulation of VAEs, the prior and the posterior distributions over the latent variables are assumed to
Third workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (ICML 2021).
Copyright 2021 by the author(s).
be both Gaussian. However, these assumptions are often incompatible and are thus limiting the performance of VAEs for complex modelling tasks (Tomczak & Welling, 2018;
Chen et al., 2018).
A natural solution to this problem is to parameterize the prior, sometimes also the posterior, with more expressive distributions.
In this preliminary work, we improve VAEs by demonstrating how DDPMs can be used for modelling the prior distribution of the latent variables.
In addition to boosting DDPM with the compression properties of VAEs, combining the two models should eventually lead to greater generative performance by enabling complex generative modelling even with simple decoder architecture.
Finally, working in the latent space should eventually reduce the computational burden associated with diffusion generative models.
2. Latent generative models 2.1. Variational autoencoder
We want to learn a generative model of an unknown distribution p ( x ) given a dataset X ∈ R N × d of N i.i.d observations x sampled from this unknown distribution.
The original VAE postulates a two-step generative process in which some unobserved variables z ∈ R h are ﬁrst sampled from a prior distribution p ( z ) and then observations x are generated from a conditional distribution p θ ( x | z ) .
The generative process can be expressed mathematically as z ∼ p ( z ) and x ∼ p θ ( x | z ) .
[Equation]
The prior p ( z ) is chosen Gaussian while the likelihood p θ ( x | z ) is modeled with a neural network.
The likelihood model decodes latent variables into observations and is thus usually refereed as the decoder in the literature.
In its original formulation, the likelihood is parameterized with a multivariate Gaussian N ( µ θ ( z ) , diag( σ 2 θ ( z ))) when the observations are continuous, and a categorical distribution when they are discrete.
Training the generative model is achieved by ﬁnding the parameters θ of the decoder that maximize the sum of the marginal likelihoods of individual points
[Equation]
These integrals are intractable but the introduction of an encoder network that approximates the posterior distribution q φ ( z | x ) allows maximizing the associated evidence lower bound
[Equation]
The ELBO becomes tighter as the approximate posterior q ψ ( z | x ) gets closer to the true posterior.
Learning the generative model is ﬁnally performed by jointly optimizing the parameters θ of the decoder and φ of the approximate posterior via stochastic gradient ascent.
In the original VAE, the encoder models the approximate posterior as a conditional multivariate Gaussian distribution N ( µ φ ( x ) , diag( σ 2 φ ( x ))) .
The ELBO loss presents two antagonistic goals to the encoder.
It should be able to both encodes the data accurately while being as close as possible to the prior distribution.
Consequently, the Gaussian assumptions made on both the prior and the posterior distributions are often incompatible and limit the generative performance.
A possible solution consists in learning a prior distribution that is compatible with the learned posteriors.
For example, Habibian et al.
(2019) and Chen et al.
(2017) respectively showed that autoregressive models and normalizing ﬂows (Rezende & Mohamed, 2015, NFs) greatly improve the performance of
VAEs when used as prior distributions.
In the following we present how denoising diffusion probabilistic models can be used to improve the performance of classical VAEs.
2.2. Denoising diffusion probabilistic models
Inspired by non-equilibrium statistical physics, SohlDickstein et al.
(2015) originally introduced DDPMs while
Ho et al.
(2020) demonstrated only more recently how to train these models for image synthesis, achieving results close to the state-of-the-art on this task.
DDPMs formulate generative modelling as the reverse operation of diffusion, a physical process which progressively destroys information.
Formally, the reverse process is a latent variable model of the form
[Equation]
where x 0 :
= x denotes the observations and x 1 , . . .
, x T denote latent variables of the same dimensionality as x 0 .
The joint distribution p φ ( x 0: T ) is modelled as a ﬁrst order
Markov chain with Gaussian transitions, that is
[Equation]
Similar to VAEs, the reverse Markov chain is trained on an
ELBO .
However, the approximate posterior q ( x 1: T | x 0 ) is ﬁxed to a diffusion process that is also a ﬁrst order Markov chain with Gaussian transitions,
[Equation]
where β 1 , . .
. , β T are the variance schedule that is either ﬁxed as training hyper-parameters or learned.
The ELBO is then given by
[Equation]
Provided that the variance schedule β t is small and that the number of timesteps T is large enough, the Gaussian assumptions on the generative process p φ are reasonable.
Ho et al.
(2020) take advantage of the Gaussian transitions to express the ELBO as
[Equation]
The inner sum in Equation (11) is made of comparisons between the Gaussian generative transitions p φ ( x t − 1 | x t ) and the conditional forward posterior q ( x t − 1 | x t , x 0 ) which can also be expressed in closed form as Gaussians
N (˜ µ t ( x 0 , x t ) , ˜ β t I ) , where ˜ β t are functions of the variance schedule.
The KL can thus be calculated with closed form expressions which reduces the variance of the ﬁnal expression.
In addition, Ho et al.
(2020) empirically demonstrate that it is sufﬁcient to take optimization steps on uniformly sampled terms of the sum instead of computing it completely.
The ﬁnal objective closely resembles denoising score matching over multiple noise levels (Song & Ermon, 2019).
These observations combined with additional simpliﬁcations leads to a simpliﬁed loss
[Equation]
where ˜ µ t ( x 0 , x t ) is the mean of q ( x t − 1 | x 0 , x t ) , the forward diffusion posterior conditioned on the observation x 0 .
3. Prior modelling with denoising diffusion
We now introduce our contribution which consists in using a DDPM for modelling the prior distribution in VAEs.
We formulate the generative model as
[Equation]
where φ denotes the parameters of the reverse diffusion model encoding the prior distribution.
Equations (13) and (14) implicitly deﬁne a prior distribution over the usual latent variables z 0 which is modelled with a reverse diffusion process.
Unfortunately, we cannot train a VAE with a diffusion prior directly on the ELBO as expressed in Equation (2) as p φ ( z 0 ) cannot be evaluated.
However, Equation (2) can be further developed as
[Equation]
in which a lower bound on the last term can be expressed by
Equation (10).
This ﬁnally leads to the following expression
E
[Equation]
≤
[Equation]
p
[Equation]
[Equation]
19) which is a valid ELBO.
Finally [Equation] the diffusion prior p φ is trained jointly with the approximate posterior q ψ and the likelihood models p θ which are optimized as in a classical
VAE.
This leads to the following loss function:
[Equation]
4. Related work
Various approaches have been proposed to improve the modelling capacity and the training of VAEs.
As a ﬁrst example, some state-of-the-art deep generative models based on VAEs model the posterior with normalizing ﬂows or autoregressive models (Kingma et al., 2016; Vahdat & Kautz, 2020).
Autoregressive models are also often used as a replacement of the original likelihood parameterization, which assumes conditional independencies that are often unrealistic (Oord et al., 2016).
Another popular improvement made to the original VAE is the embedding of structure in the latent variables.
In particular, hierarchical VAEs (Sønderby et al., 2016; Kingma et al., 2016) combined with careful training demonstrate impressive results on generative modelling for images (Vahdat & Kautz, 2020). Vahdat et al.
(2021) concurrently proposed to use diffusion for modelling the prior distributions of VAEs.
They obtain state-of-the-art results on image synthesis by combining continuous diffusion models and VAEs.
Not as close to our work but related, Chen et al.
(2017) proposed to learn the prior as a solution to the mismatch between the approximate and the true posteriors.
They model the prior with an autoregressive ﬂow, which also closely relates to modelling the posterior distribution with an inverse autoregressive ﬂow (Kingma et al., 2016).
Tomczak & Welling (2018) takes inspiration from the aggregated posterior 1 N (cid:80) Ni =
1 q ψ ( z | x ) (Hoffman & Johnson, 2016; Makhzani et al., 2015) to introduce the VampPrior deﬁned as a mixture of learned pseudoinputs.
An orthogonal line of work suggests that the mismatch between the approximate posterior and the exact posterior can be reduced by over-weighting the terms related to the prior and to the approximate posterior in the ELBO (Higgins et al., 2016;
Chen et al., 2018).
5. Experiments
We now compare VAEs for different choices of priors, including the original Gaussian prior, an NF prior, and the proposed diffusion prior.
All models share a same backbone encoder-decoder architecture inspired from DCGAN (Radford et al., 2015).
Optimization is performed with Adam for 250 epochs with a learning rate set to 0 . 0005 .
After each epoch, the models are evaluated on a validation set used to select the best one for each training setting.
We compare the models on the CIFAR10 and CelebA datasets for 3 different latent variables dimensionality ( 40 , 100 , 200 ).
The NF used in our experiments is a 3-step autoregressive afﬁne ﬂow with simple MLP backbones similar to the one used to model the transition function of DDPM.
Table 1 presents the FID scores for the different models.
We ﬁrst notice the large scores reached by all models on the
CIFAR10 dataset.
This can be explained by the simplicity of the models trained in our experiments.
We believe these scores could be greatly improved by using a more sophisticated likelihood model such as a PixelCNN (Oord et al., 2016).
Although FID scores suggest that the Gaussian prior outperforms the diffusion prior in terms of generative performance, the visual inspection of Figure 1 shows that the diffusion prior results in samples slightly more realistic than those of the classical VAE.
The best FID score is achieved by the NF prior, although its samples do not seem to reﬂect this superiority.
In this case, we believe the FID scores are not entirely informative about the quality of the images synthesized by the models and should be interpreted with a grain of salt.
Although learned priors seem to improve generative performance on CIFAR10, additional work is needed to reach results that would justify using a diffusion prior for this dataset.
On CelebA however, we observe in Table 1 that diffusion priors outperform the Gaussian prior.
This is in line with [Table] the visual inspection of Figure 2a and Figure 2c.
As for
CIFAR10, the NF prior outperforms the Gaussian and diffusion priors in terms of FID scores, although the visual inspection of the corresponding samples in Figure 2b does not reveal a much better quality of images when compared to those resulting from the diffusion prior.
We conclude from these observations that diffusion priors offer an interesting alternative to NFs for modelling the prior in a VAE.
6. Conclusion and future work
This preliminary work presents how denoising diffusion probabilistic models can be used as a new class of learnable priors for VAEs.
As a notable contribution, we empirically demonstrate that optimizing implicitly a prior on an ELBO can be performed jointly to training the encoder and the decoder of the VAE.
In addition, our results suggest DDPM performs on par with NFs for modelling prior distribution.
A large spectrum of future research directions could beneﬁt from the basic idea expressed in this preliminary work.
As an example, recent advances in diffusion models such as the continuous formulation (Song et al., 2020) or improvement to the training procedure of DDPM (Nichol & Dhariwal, 2021) could be implemented in the prior model.
Similarly, many improvements could be made to the architectures used for the VAE and to the training procedure.
In particular, image synthesis with hierarchical VAEs which organizes the latent variables into multiple scales images could reveal the full potential of diffusion priors.
This would indeed combine the structural knowledge embed by such type of VAEs with the impressive performance of DDPM for modelling distributions over images.
Finally, diffusion does not constrain the neural networks architectures and so enables the embedding of a larger choice of inductive biases in the prior distribution compared to autoregressive models and NFs.