VAEs meet Diffusion Models: Efﬁcient and High-Fidelity Generation
Abstract
Diffusion Probabilistic models have been shown to generate state-of-the-art results on several competitive image synthesis benchmarks but lack a low-dimensional, interpretable latent space, and are slow at generation.
On the other hand, Variational Autoencoders (VAEs) have access to a low-dimensional latent space but, despite recent advances, exhibit poor sample quality.
We present VAEDM, a novel generative framework for reﬁning VAE generated samples using diffusion models while also presenting a novel conditional forward process parameterization for diffusion models.
We show that the resulting parameterization can improve upon the unconditional diffusion model in terms of sampling efﬁciency during inference while also equipping diffusion models with the low-dimensional VAE inferred latent code.
Furthermore, we show that the proposed model exhibits out-of-the-box capabilities for downstream tasks like image superresolution and denoising.
Introduction
Generative modeling is the task of capturing the underlying data distribution and learning to generate novel samples from a posited explicit/implicit distribution of the data in an unsupervised manner.
Variational Autoencoders (VAEs) [17, 29] are a type of explicit-likelihood based generative models which can also learn a low-dimensional latent representation for the data.
The resulting framework is ﬂexible and can be used for several downstream applications [3, 11, 18, 27].
However, in image synthesis applications, VAE generated samples are usually blurry and fail to incorporate highfrequency information [7].
Despite recent advances [4, 28, 37, 39] in improving VAE sample quality, there is still a signiﬁcant gap in sample quality between VAEs and their implicit-likelihood counterparts like GANs [8, 14, 15, 16].
In contrast, Denoising diffusion probabilistic models (DDPM)
[12, 33] have been recently shown to achieve impressive performance on several image synthesis benchmarks, even surpassing GANs on several such benchmarks [6, 13].
However, conventional diffusion models require an expensive iterative sampling procedure and lack a low-dimensional latent representation, limiting these models’ practical applicability for downstream applications.
In this preliminary work, we present VAEDM, a novel framework which combines the best of both
VAEs and DDPMs.
VAEDM consists of a generator-reﬁner framework in which blurry samples generated from a VAE are reﬁned using a conditional DDPM.
Using qualitative and quantitative experiments, we show that the proposed method, unlike standard VAEs, can generate high-quality samples that can be controlled using the low-dimensional VAE latent space.
Moreover, we show that the proposed method requires fewer reverse diffusion sampling steps during inference, and exhibits out-of-the-box generalization to downstream tasks like image super-resolution and denoising.
The VAEDM framework
Before introducing the VAEDM framework, we recommend the readers to refer to Appendix A for a background on VAEs and diffusion models.
Given a high-resolution image x 0 , an auxiliary conditioning signal y to be modelled using a VAE, a latent representation z associated with y , and a sequence of T representations x 1: T learned by a diffusion model, the VAEDM generative process can be formulated as follows:
[Equation]
where θ and φ are the parameters of the VAE decoder and the reverse process of the conditional diffusion model, respectively.
Since computation of the likelihood for this generative process is intractable, we can approximate it by computing a lower bound (ELBO) with respect to the joint posterior over the unknowns ( x 1: T , y , z ) which can be formulated as follows:
[Equation]
where ψ are the parameters of the VAE recognition network ( q ψ ( z | y, x 0 ) ).
We keep the DDPM forward process ( q ( x 1: T | y, z, x 0 ) ) ﬁxed throughout training.
It can be shown that the speciﬁed probabilistic framework yields the following ELBO objective (See Appendix C.1 for proof)
[Equation]
Therefore, the overall VAEDM training objective decomposes into a sum of VAE and a conditional
DDPM objectives.
In addition to the above objective, we make the following simplifying assumptions:
We assume the conditioning signal y to be x 0 itself.
Given this choice, we do not condition the reverse diffusion process on y and take it as p φ ( x 0: T | z ) . For ease of optimization, we train Eq.
(3) in a sequential two-stage manner, i
e .,
ﬁrst optimizing L VAE and then optimize for L DDPM while ﬁxing θ and ψ . 3.
Lastly, instead of conditioning the reverse diffusion directly on the latent code z , we condition the second stage DDPM model on the VAE reconstruction ˆ x 0 which is a deterministic function of z .
Moreover, in this work, we only consider a single-stage VAE (with a single stochastic layer) for the
VAE training stage. However, due to the ﬂexibility of the VAEDM two-stage training, multi-stage VAE approaches as proposed in [4, 28, 37] can also be utilized.
For the second stage DDPM training, we consider the following two DDPM variants in this work.
Formulation-1 :
In this DDPM formulation, we assume that the forward process is independent of the VAE reconstructions ˆ x (and hence the latent code information z ) i.e. q ( x 1: T | z, x 0 )
≈ q ( x 1: T | x 0 ) .
Moreover, given the simplifying assumptions discussed previously, the reverse process transitions take the following form i.e. p ( x 0: T | z ) ≈ p ( x 0: T | ˆ x 0 ) .
Formulation-2 :
In this DDPM formulation, we assume that the forward process is also dependent on the VAE reconstructions ˆ x i.e. q ( x 1: T | z, x 0 )
≈ q ( x 1: T | x 0 , ˆ x 0 ) .
The form of the reverse process remains the same as in Formulation-1.
Since, the forward process is also dependent on ˆ x 0 in addition to the input data x 0 , the form of the conditional marginal q ( x t | x 0 ) in the forward process changes, which leads to a modiﬁed DDPM training and inference procedure.
We refer interested readers to
Appendix B for a detailed discussion of this formulation.
Experiments
We evaluate the effectiveness of our method by making qualitative and quantitative comparisons between the proposed VAEDM framework and the unconditional DDPM model [12] on the
[Figure]
VAE samples
VAEDM (Formulation-1)
VAEDM (Formulation-2)
VAE Sample(Fixed)
VAEDM FormulationVAEDM Formulation- CelebAMask-HQ [22] dataset. All the images in the CelebAMask-HQ dataset were downsampled to 128 x 128 resolution for efﬁciency.
For all the experiments, we set number of DDPM timesteps, T = 1000 , the noise schedule, β 1 = 10 − 4 and β 2 = 0 .
02 and use our unconditional DDPM implementation to compare with the proposed approach. More details regarding the network architecture and the training and evaluation procedures can be found in Appendix E.
3.1 Sample quality
Fig.
1a shows some samples generated from the VAEDM model (using Formulation-1).
As can be observed from the visualization, the ﬁnal generated samples (Fig. 1a (Top row)) are a reﬁnement of the blurry samples generated by our single-stage VAE model (Fig.
1a (Bottom row)).
Some more unconditional samples from the VAEDM model are provided in Appendix F.
3.2 Number of Reverse process steps during sampling
Next, we qualitatively compared the unconditional DDPM model proposed in [12] and VAEDM in terms of the number of reverse process sampling steps required to produce high-ﬁdelity samples.
The visualization in Fig.
1b suggests that our method requires around half the number of reverse process steps to generate high-ﬁdelity samples than the unconditional DDPM model.
These qualitative results are further supported by Fig. 3a where VAEDM largely outperforms the unconditional DDPM model in terms of FID in the low-step regime.
Intuitively, these results might not be surprising as the DDPM model in VAEDM only needs to reﬁne a pre-generated blurry template while the unconditional
DDPM in [12] requires a longer sampling schedule to model all image details, which is a harder task than reﬁning a pre-generated template.
3.3 Interpolations in the latent space
The proposed VAEDM model consists of two latent codes: the VAE latent code z and the latent representations x T associated with the reverse process base measure p ( x T ) (which is of the same size of the input image x 0 and thus might not be beneﬁcial for downstream tasks).
We next discuss the effects of manipulating both z and x T .
We consider the following two interpolation settings:
Varying z with ﬁxed x T : We ﬁrst sample two VAE latent codes z 1 and z 2 using the standard Gaussian distribution.
We then perform linear interpolation between z 1 and z 2 to obtain intermediate VAE latent codes ˜
z = λz 1 + (1 − λ )
z 2 , which are then used to generate the ﬁnal VAEDM samples with a shared x T latent.
Fixed z with varying x T : Next, we sampled the VAE latent code z using the standard Gaussian distribution.
With a ﬁxed z , we then sampled two latent DDPM representations x T 1 and x T 2 from the [Figure] D [Figure]
Noisy Images RestoredImages OriginalImages 16 x32 x Noise std =0.
Image Superres
Image Denoising
[Figure]
reverse process base measure p ( x T ) .
We then performed linear interpolation between x T 1 and x Twith a ﬁxed z to generate the ﬁnal VAEDM samples.
As can be observed from Fig. 2, interpolating in the low-dimensional VAE latent space leads to changes in major features of the generated samples while changes in the DDPM latent conditioned on the same VAE reconstructions leads to minor differences between the generated samples which can be attributed to the stochastic nature of the DDPM sampling procedure.
This observation implies that we can control the DDPM generated samples primarily by manipulating the low-dimensional VAE latent code z vae and that the DDPM latents in VAEDM have low-entropy.
To the best of our knowledge, our work is the ﬁrst to equip diffusion models with such a low-dimensional latent representation.
3.4 Out-of-the-box super-resolution and denoising
Recently, [31] showed impressive image super-resolution results using diffusion models.
To test if our generator-reﬁner framework can generalize over different types of noisy inputs, we conditioned the DDPM reverse process on the noisy input (instead of the VAE reconstruction).
The samples obtained from such conditioning are visualized in Fig.
On a 16 x 16 to 32 x 32 image super-resolution task and an image denoising task (using gaussian noise corruption) for the CIFAR-10 dataset, our model is able to generate plausible reconstructions (Fig.
3b(Middle row))
when compared to the original samples. Intuitively, these results can be expected since the task of reﬁning (blurry) VAE reconstructions might be more challenging than learning to upsample a downsampled version of an image.
However, it is worth noting that certain artifacts in the generated reconstructions are evident, leaving scope for improvements.
Conclusion
In this work, we presented a novel unifying framework for training VAEs and diffusion models.
We presented the effectiveness of the proposed approach in generating high-quality samples, requiring fewer reverse process steps during inference when compared with the unconditional DDPM formulation, equipping DDPM with a low dimensional latent code, and generalizing to additional tasks like image super-resolution and denoising.
We also provide a detailed comparison of our proposed approach with other state-of-the-art generative model families in Appendix D.
However, we look forward to combining our proposed approach with recent advances in diffusion models and make more quantitative comparisons with other state-of-the-art methods on image-synthesis benchmarks.
Impact statement
We note that the synthetic image generation techniques have the potential to mitigate bias and privacy issues for related ML models that require data collection and annotation.
However, such techniques could be misused to produce fake or misleading information, and researchers should be aware of these risks and explore the techniques responsibly.
Acknowledgements
We would like to thank Ben Poole for his insightful comments during the course of this project
A Appendix: Background on VAE’s and Diffusion models A.1 Variational Autoencoders
VAEs [17] are based on a simple but principled encoder-decoder based formulation which tries to maximize the evidence lower bound (ELBO) of the data log-likelihood, which is intractable to compute in general.
The VAE optimization objective can be stated as follows:
[Equation]
Under the amortized variational inference scheme, the approximate posterior ( q φ ( z | x ) ) and the likelihood ( p θ ( x | z ) ) distributions can be modeled using deep neural networks with parameters φ and θ , respectively using the reparameterization trick [17, 29].
The choice of the prior distribution p ( z ) is ﬂexible and can vary from a standard Gaussian [17] to more expressive priors like normalizing ﬂows [9, 19, 38].
A.2 Denoising Diffusion Probabilistic Models
DDPMs [12, 33] are latent-variable models consisting of a forward noising process ( q ( x 1: T | x 0 ) ) which gradually destroys the structure of the data x 0 and a reverse process (( p ( x 0: T ) )) which learns to recover the original data x 0 from the noisy input.
The forward noising process is modeled using a ﬁrst-order Markov chain with Gaussian transitions and is ﬁxed throughout training, and the noise schedules β 1 to β T can be ﬁxed or learned.
The form of the forward process and some notable properties of the forward process conditional distributions are summarized in the equations below
[Equation] [Equation]
The forward process of DDPMs admits a closed form for x t for any t , as follows
[Equation]
¯ α t x 0 , (1 − ¯ α t ) I ) where α t =
(1 − β t ) and ¯ α t =
(cid:89) t
[Equation]
t
[Equation]
[Equation]
The forward process posteriors are also tractable and are given by
[Equation]
where
[Equation]
The reverse process can also be parameterized using a ﬁrst-order Markov chain with a learned
Gaussian transition distribution as follows
[Equation]
Given a large enough T and a well-behaved variance schedule of β t , the distribution q ( x T | x 0 ) will approximate an isotropic Gaussian.
We can generate a new sample from the underlying data distribution q ( x 0 ) by sampling a latent from p ( x T ) (chosen to be an isotropic Gaussian distribution) and running the reverse process.
As proposed in [12], the reverse process in DDPM is trained to minimize the following upper bound over the negative log-likelihood (See [33] for detailed proofs):
[Equation]
A notable aspect of the above objective is that all the KL divergences involve Gaussians and, consequently, are available in closed form.
Notably, [12] parameterize the reverse process conditional p θ ( x t − 1 | x t ) using the forward process posterior q ( x t − 1 | x t , x 0 ) .
[12] show that such a parameterization simpliﬁes the second term in Eq.
10 at any given time-step t to the following objective in Eq. [Equation] where x t = √ ¯ α t x 0 + (cid:15) √ 1 − ¯ α t and (cid:15) ∼ N (0 , I ) . Intuitively, this means that the reverse process in
DDPM is trained to predict the noise added to the input x 0 at any time-step t .
We use this simpliﬁed training formulation throughout our work to train all proposed parameterizations of diffusion models as [12] show that this formulation yields superior sample quality than other forms of reverse process parameterizations.
For further details on the exact training and inference processes, we encourage the readers to refer to [12].
B Appendix: DDPM training and inference under Formulation
[Table]
[Equation]
[Table]
The DDPM training objective proposed in [12], has the following form:
[Equation]
Reverse Process parameterization:
Following [12], we parameterize the reverse process transition p θ ( x t − 1 | x t ) using the functional form of the forward process posterior q ( x t − 1 | x t , x 0 ) .
For the
VAEDM formulation-2 proposed in Section 2, we design the forward process conditional distributions as follows:
[Equation]
Given this choice of q ( x t | x t − 1 , ˆ x 0 ) , it can be shown that the conditional marginal distribution q ( x t | x 0 , ˆ x 0 ) can be speciﬁed as:
[Equation]
The posterior distribution q ( x t − 1 | x t , x 0 ) will also be a Gaussian distribution with the following form:
[Equation]
where ˆ µ t ( x t , x 0 , ˆ x 0 ) =
β t
[Equation]
ˆ β t =
(1 − ¯ α t − 1 ) 1 − ¯ α t β t and x 0 = 1 √ ¯ α t (
[Equation]
√ 1 − ¯ α t ) where (cid:15) ∼ N (0 ,
[Equation]
[Equation]
Hence the forward process posterior in the VAEDM formulation is a shifted version of the forward process posterior proposed in [12].
Since the VAE reconstruction ˆ x 0 for an image x 0 is constant during DDPM training, we can parameterize the reverse process posterior as ˆ µ θ ( x t , x 0 , ˆ x 0 , t )
=
˜ µ θ ( x t , x 0 , t )+ κ ˆ x 0 .
Additionally, we keep the variance of the reverse process conditional ﬁxed and equal to ˆ β t as proposed in [12].
Since L t − 1 ∝ (cid:107) ˆ µ t ( x t , x 0 , ˆ x 0 ) − ˆ µ θ ( x t , x 0 , ˆ x 0 , t ) (cid:107) 2 , the DDPM training objective in our formulation remains unchanged from the simpliﬁed denoising score matching objective proposed in [12].
Choice of the decoder, L 0 : One possible choice for the decoder is to set p θ ( x 0 | x 1 ) to be a discrete independent decoder derived from the Gaussian N (ˆ µ θ ( x 1 , ˆ x 0 , 1) , ˆ β 1 I ) [12].
However, at t = 1 , we have ˆ µ θ ( x 1 , ˆ x 0 , 1) =
x 0 ( x 1 , ˆ x 0 , (cid:15) θ ) + ˆ x 0 .
Therefore, to account for the VAE reconstruction bias in the ﬁnal DDPM output, we set our decoder p θ ( x 0 | x 1 ) =
N (ˆ µ θ ( x 1 , ˆ x 0 , 1) − ˆ x 0 , ˆ β 1 I ) .
Without using this adjustment, we found the ﬁnal DDPM samples to be a bit blurry in our initial experiments.
The ﬁnal training and inference algorithms are summarized in Algorithms 1 and 2 respectively.
In Algorithm 2, the coefﬁcients γ 0 , γ 1 and γ 2 denote the coefﬁcients of the forward process posterior in Eqn. 17.
C Appendix: Detailed Proofs C.1 Derivation of VAEDM ELBO objective
Given a high-resolution image x 0 , an auxiliary conditioning signal y to be modelled using a VAE, a latent representation z associated with y , and a sequence of T representations x 1: T learned by a diffusion model, the DiffuseVAE generative process, p ( x 0: T , y, z ) can be factorized as follows:
[Equation]
where θ and φ are the parameters of the VAE decoder and the reverse process of the conditional diffusion model, respectively.
The log-likelihood of the training data can then be obtained as:
[Equation]
Furthermore, since the joint posterior p ( x 1: T , z | y, x 0 ) is intractable to compute, we approximate it using a surrogate posterior q ( x 1: T , z | y, x 0 ) which can also be factorized into the following conditional distributions:
[Equation]
where ψ are the parameters of the VAE recognition network ( q ψ ( z | y, x 0 ) ).
Since computation of the likelihood in Eq. (20) is intractable, we can approximate it by computing a lower bound (ELBO) with respect to the joint posterior over the unknowns ( x 1: T , z ) as:
[Equation]
Plugging the factorial forms of the DiffuseVAE generative process and the joint posterior deﬁned above in eqn. (22), we can simplify the ELBO as follows: log p ( x 0 , y ) ≥ E q ( x 1: T ,z | y,x 0 ) (cid:20) log p ( x 0: T , y, z ) q ( x 1: T , z | y, x 0 ) (cid:21) (23) ≥ E q ( x 1: T ,z | x 0 ,y ) (cid:20) log p ( z ) p θ ( y | z ) p φ ( x 0: T | y, z ) q ψ ( z | y, x 0 ) q ( x 1: T | y, z, x 0 ) (cid:21) (24) ≥ E q ( x 1: T ,z | x 0 ,y ) (cid:34) log p ( z ) q ψ ( z | y, x 0 ) + log p θ ( y | z )+ log p φ ( x 0: T | y, z ) q ( x 1: T | y, z, x 0 ) (cid:35) (25) ≥ E q ( z | y,x 0 ) (cid:20) log p ( z ) q ψ ( z | y, x 0 ) + log p θ ( y | z ) (cid:21) + E q ( x 1: T ,z | x 0 ,y ) (cid:20) log p φ ( x 0: T | y, z ) q ( x 1: T | y, z, x 0 ) (cid:21) (26) ≥ E z ∼ q ( z | y,x 0 ) (cid:34)
E q ( x 1: T | y,z,x 0 ) (cid:20) p φ ( x 0: T | y, z ) q ( x 1: T | y, z, x 0 ) (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) L DDPM (cid:35) (27) + E q ψ ( z | y,x 0 ) [ p θ ( y | z )] − D KL ( q ψ ( z | y, x 0 ) || p ( z ))
[Equation]
C.2 Derivation of the forward conditional marginal in Formulation
Given the forward process transitions for VAEDM (Formulation-2)(See Appendix B): [Equation]
[Equation]
From Eqn.(30)
[Equation]
we can write,
[Equation]
where
[Equation]
Taking expectations both sides,
[Equation]
Substituting E ( x 1 ) = √ 1 − β 1 x 0 + ˆ x 0 from Eqn.(29) into the above formulation we get,
[Equation]
Similarly it can be shown that V ar ( x t ) =
(1 − ¯ α t ) I .
Therefore,
[Equation]
D Appendix: Detailed Comparison with existing approaches
Following the seminal work of [12, 33] in diffusion models, there has been a lot of recent progress in both unconditional [6, 20, 26] and conditional diffusion models [2, 5, 13, 31] (including score-based models [34, 35], based on a connection proposed in [12]) for a variety of downstream tasks including image synthesis, audio synthesis and likelihood estimation among others.
Similarly there has also been progress in improving the ELBO estimates [1, 24, 32] and image synthesis [4, 23, 37, 40] quality using VAE’s [17, 29].
Next, we compare our proposed approach in detail with several of these related existing model families.
Unconditional DDPM :
DDPM as introduced in [12] generates images unconditionally which has limited application scope.
On the other hand, the proposed VAEDM model can be used for tasks like image enhancement, super-resolution etc.
Essentially, the DDPM formulations proposed in our method can also be trained in a stand-alone manner with the reverse process conditioned on an auxiliary input like a grayscale image or an image with missing pixels to generate colorized or inpainted results, respectively.
Moreover, our approach requires lesser number of reverse process steps in general to produce plausible samples.
Conditional DDPM :
Conditional DDPM as introduced in [13] and [31] use multiple diffusion models for generating high-resolution images in a cascaded fashion.
However, for even a two-stage pipeline the sampling time of such models would be effectively much higher than VAEDM.
Given the ﬂexibility of our approach, we hypothesize that a single-stage VAE can also be replaced by a complex multi-stage VAE architecture as proposed in [4, 37] for comparable sample quality to cascaded diffusion models without affecting the sampling time signiﬁcantly.
Moreover, such cascades lack a low-dimensional latent code which might be a limiting factor for certain downstream applications.
It is worth noting that, [13] use a conditioning augmentation scheme where the high resolution image is generated by conditioning on a blurred/noisy low resolution image.
This augmentation scheme is empirically chosen and can be sub-optimal across datasets.
On the contrary, our model is conditioned on a reconstruction generated by a VAE and no explicit augmentation is required.
Hierarchical VAEs Hierarchical VAEs [4, 28, 36, 37] can suffer from posterior collapse and heuristics like gradient skipping and spectral normalization [25] might be required to stabilize training.
Moreover, these models require a large dimensionality of the latent codes to generate high-ﬁdelity samples [28, 39].
In contrast, VAEDM training does not suffer from such instabilities and the proposed method requires a single latent code layer (with dimensionality comparable to GANs) to generate high-ﬁdelity samples.
Interestingly, in principle, our approach can be considered similar to
Hierarchical VAEs in some sense since the ﬁrst stage VAE is used to model the high-level attributes of the generated image while the DDPM model in the second stage can be considered as a reﬁner that adds low-level details.
In Hierarchical VAEs, the subsequent latent hierarchies perform this function.
E Model Architecture and training
Unless otherwise mentioned, we re-used the model architectures and training hyperparameters proposed in [12].
For all experiments, we used the EMA approach from [10] for updating the target model parameters with an initial decay rate of 0.9999.
However, we did not experiment with the choice of the EMA scheme.
Data preprocessing : For experiments with the CelebAMask-HQ dataset [22], we downsampled the training data to 128 x 128 resolution since we believe this setting provides the correct balance between visual details and training efﬁciency.
We did not use any form of data augmentation during training for both CelebAMask-HQ and CIFAR-10 datasets.
For Stage-1 VAE training, the training data was scaled between [0.0, 1.0].
For Stage-2 DDPM training, the images were scaled between (both formulations).
Model architecture :
The VAE architecture used for Stage-1 training on the CelebAMask-HQ dataset consists of around 21M parameters with residual block architectures inspired from [4].
The VAE latent code size was set to 1024 for CelebAMask-HQ and 512 for the CIFAR-10 dataset.
The U-Net [30] decoder used in Stage-2 DDPM training is around 113M parameters.
For CIFAR-10, the VAE model consists of around 9M parameters whereas the size of the U-Net decoder used is 34.8M parameters.
Training : The VAE model was trained for around 180 epochs and 500 epochs for the CelebAMaskHQ and the CIFAR-10 datasets with batch sizes of 40 and 128 respectively.
All DDPM models were trained for around 250 epochs with a batch size of 64 for the CelebAMask-HQ dataset and for around 1000 epochs with a batch size of 128 for the CIFAR-10 dataset.
We used a mix of 4 Nvidia 1080Ti GPUs ( 44GB memory) and a cloud TPUv2-8 ( 64GB memory) provided by Colab Pro for training the models.
For DDPM training, it took around 0.74s and 0.2s per optimization step for the
CelebAHQ and CIFAR-10 datasets respectively.
F Unconditional samples
using VAEDM
[Figure] [Figure]