# Diffusion Variational Inference: Diffusion Models as Expressive Variational
Posteriors

Top Piriyakulkij
Cornell University
&Yingheng Wang
Cornell University
&Volodymyr Kuleshov
Cornell Tech

###### Abstract

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.

##  1 Introduction

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.

This work seeks to improve variational inference via expressive posteriors
based on diffusion models (Ho et al., 2020; Song et al., 2020). Diffusion
methods are defined via a noising process, which maps data into Gaussian
noise; a diffusion model generates data by reversing this noising process,
which yields high-quality samples and accurate density estimates (Kingma et
al., 2021). Here, we argue for using diffusion models in latent space, where
we gradually map a simple (e.g., Gaussian) latent representation of the data
into one that is more complex via an iterative diffusion-like procedure. This
procedure yields an expressive approximate posterior trained with a denoising
objective that does not involve adversarial training (Makhzani et al., 2015)
or constrained invertible architectures (Kingma et al., 2016).

Specifically, we propose denoising diffusion variational inference (DDVI), an
approximate inference algorithm that introduces auxiliary latent variables
into the approximate posterior via a user-specified noising process. This
process transforms the latent variable we seek to model into a simple (e.g.,
Gaussian) auxiliary latent; during inference, we fit the approximate posterior
by reversing this noising process. Our learning objective is a variational
lower bound inspired by the wake-sleep algorithm (Hinton et al., 1995) that
can be interpreted as a form of regularized variational inference. We also
derive extensions of our method to semi-supervised learning and clustering.

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 evaluate DD-VAEs on
a real problem in biological data analysis—inferring human ancestry from
genetic data. Our method outperforms strong baselines on the Thousand Genomes
dataset (Siva, 2008) and learns a low-dimensional latent space that preserves
semantically meaningful structure (Haghverdi et al., 2015).

##### Contributions.

In summary, this work introduces denoising diffusion variational inference, an
approximate inference algorithm that features three key components: auxiliary
latent variables, a user-specified noising process over these variables, and a
lower bound on the marginal likelihood inspired by wake-sleep. Our method can
be used to fit deep latent variable models, which yields the DD-VAE algorithm.
This algorithm is especially effective at dimensionality reduciton and
representation learning, where it outperforms alternative methods based on
adversarial training.

Figure 1: Denoising diffusion variational inference in a VAE. Between the
encoder and decoder, we have a diffusion model to map a simple distribution
into a complex distribution over latents.

##  2 Background

##### Deep Latent Variable Models

Latent variable models (LVMs) $p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$ are
usually fit by optimizing the evidence lower bound (ELBO)

$\log
p_{\bm{\theta}}(\mathbf{x})\leq\mathbb{E}_{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})]-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{z}))$

which serves as a tractable surrogate for the marginal log-likelihood (MLL).
The gap between the MLL and the ELBO equals precisely
$D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{z}|\mathbf{x}))$—thus,
a more expressive $q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})$ may better fit the
true posterior and induce a tighter ELBO (Kingma & Welling, 2013).

Expressive variational posteriors can be formed by choosing more expressive
model families—including auxiliary variable methods (Maaløe et al., 2016),
MCMC-based methods (Salimans et al., 2015), normalizing flows (Rezende &
Mohamed, 2015)—or improved learning objectives—e.g., adversarial or sample-
based losses (Makhzani et al., 2015; Zhao et al., 2017; Si et al., 2022;
2023).

The wake-sleep algorithm (Hinton et al., 1995) optimizes an alternative
objective

$\mathbb{E}_{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})]-D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))$

, in which the KL divergence term is reversed. The learning procedure for
wake-sleep involves alternating between ”wake” phases where the recognition
model is updated and ”sleep” phases where the generative model is refined.

##### Denoising Diffusion Models

A diffusion model is defined via a user-specified noising process $q$ that
maps data $\mathbf{x}_{0}$ into a sequence of $T$ variables
$\mathbf{y}_{1:T}=\mathbf{y}_{1},...,\mathbf{y}_{T}$ that represent increasing
levels of corruption to $\mathbf{x}_{0}$. We obtain $\mathbf{y}_{1:T}$ by
applying a Markov chain
$q(\mathbf{y}_{1:T}|\mathbf{x}_{0})=\prod_{t=1}^{T}q(\mathbf{y}_{t}|\mathbf{y}_{t-1})$,
where we define $\mathbf{y}_{0}=\mathbf{x}_{0}$ for convenience. When
$\mathbf{x}_{0}$ is a continuous vector, a standard choice of transition
kernel is
$q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{y}_{t};\sqrt{{\bm{\alpha}}_{t}}\mathbf{y}_{t-1},\sqrt{1-{\bm{\alpha}}_{t}}\mathbf{I})$,
which is a Gaussian centered around a copy of $\mathbf{y}_{t-1}$ to which we
added noise following a schedule
$0<{\bm{\alpha}}_{1}<{\bm{\alpha}}_{2}<...<{\bm{\alpha}}_{T}=1$.

A diffusion model can then be represented as a latent variable distribution
$p(\mathbf{x}_{0},\mathbf{y}_{1:T})$ that factorizes as
$p(\mathbf{x}_{0},\mathbf{y}_{1:T})=p(\mathbf{y}_{T})\prod_{t=0}^{T-1}p_{\bm{\theta}}(\mathbf{y}_{t}\mid\mathbf{y}_{t+1})$
(again using $\mathbf{y}_{0}$ as shorthand for $\mathbf{x}_{0}$). This model
seeks to approximate the reverse of the forward diffusion $q$ and map noise
$\mathbf{y}_{T}$ into data $\mathbf{x}_{0}$.

The true reverse of the process $q$ cannot be expressed in closed form; as
such, we parameterize $p_{\bm{\theta}}$ with parameters ${\bm{\theta}}$
trained by maximizing the ELBO:

$\log p_{\bm{\theta}}(\mathbf{x}_{0})\geq\mathbb{E}_{q}\left[\log
p_{\bm{\theta}}(\mathbf{x}_{0}|\mathbf{x}_{1})-\sum_{t=2}^{T}D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0})||p_{\bm{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_{t}))\right]-D_{\mathrm{KL}}(q(\mathbf{x}_{T}|\mathbf{x}_{0})||p(\mathbf{x}_{T}))$

##### Visualization and Dimensionality Reduction

There exist two important types of dimensionality reduction methods, which (1)
emphasize the preservation of pairwise distance structures among all data
samples, e.g., PCA (Wold et al., 1987) and LDA (Balakrishnama & Ganapathiraju,
1998), and (2) local distances over global ones, e.g., t-SNE (Van der Maaten &
Hinton, 2008) and UMAP (McInnes et al., 2018). Latent variable models (LVMs)
(Kingma et al., 2016; Makhzani et al., 2015) are another class of techniques
for dimensionality reduction. They represent high-dimensional data in terms of
latent, or hidden, variables in a lower-dimensional space, effectively
providing a compact representation of the data.

##  3 Variational Inference With Denoising Diffusion Models

We introduce denoising diffusion variational inference (DDVI), which enhances
variational inference with diffusion-based methods and is motivated by
challenges in data visualization and dimensionality reduction. Our approach
consists in augmenting variational inference in a latent variable model
$p(\mathbf{x},\mathbf{z})$ with auxiliary latents $\mathbf{y}\in\mathcal{Y}$
introduced via a user-specified noising process $r(\mathbf{y}|\mathbf{z})$.
The $r(\mathbf{y}|\mathbf{z})$ transforms $\mathbf{z}$—which is the latent
whose intractable posterior we seek to approximate—into $\mathbf{y}$, whose
posterior will be easier to model. Examples of $r$ include forward diffusion
processes, discrete noising processes (Austin et al., 2021), as well as custom
regularizers (Section 3.3).

We then form an expressive posterior $q(\mathbf{z}|\mathbf{x})$ by fitting the
reverse of the noising process $r(\mathbf{y}|\mathbf{z})$ as in a diffusion
model. We define $q(\mathbf{z}|\mathbf{x})$ by sampling from a first model
$q(\mathbf{y}|\mathbf{x})$—this is an easier task since we can choose $y$ to
have a simple (e.g., Gaussian) posterior—and then by sampling from a denoising
model $q(\mathbf{z}|\mathbf{x},\mathbf{y})$ that approximates the reverse
process $r(\mathbf{z}|\mathbf{y})$. The model
$q(\mathbf{z}|\mathbf{x},\mathbf{y})$ is parameterized and trained as a
denoising diffusion model, and is thus a flexible posterior estimator.

When diffusion variational inference is used to fit a variational autoencoder,
we refer to the resulting algorithm as VAE with diffusion encoders (DD-VAE).
We define the full algorithm below.

###  3.1 DD-VAE: Variational Autoencoders With Diffusion Encoders

We seek to fit a latent variable model
$p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$ with a potentially complex prior
$p_{\bm{\theta}}(\mathbf{z})$. One perspective that can be used to define our
approach consists in lifting $p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$ into an
extended latent space by introducing auxiliary latents
$\mathbf{y}\in\mathcal{Y}$ and applying variational inference in the extended
space.

Specifically, we augment the model with the aforementioned user-specified
noising process $r(\mathbf{y}|\mathbf{z})$. We use
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})=r(\mathbf{y}|\mathbf{z})p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$
to denote the extended probability distribution. Note that the marginalizing
out $\mathbf{y}$ in $p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})$ yields
the original model: hence fitting
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})$ in an extended probability
space solves our original task.

The noising process $r$ may also introduce multiple latents
$\mathbf{y}_{1:T}$, as in the forward process
$r(\mathbf{y}_{1:T}|\mathbf{z})=\prod_{t=1}^{T-1}r(\mathbf{y}_{t+1}|\mathbf{y}_{t},\mathbf{z})$
of a diffusion model; we may then define
$r(\mathbf{y}|\mathbf{z})=r(\mathbf{y}_{T}|\mathbf{z}):=\int
r(\mathbf{y}_{T},\mathbf{y}_{1:T-1}|\mathbf{z})d\mathbf{y}_{1:T-1}$. While we
do not require this specific form for $r$, we will return to diffusion noising
processes later in the section.

####  3.1.1 Denoising Diffusion Variational Inference (DDVI)

One possible way to apply variational inference to
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})$ is to apply the ELBO twice
to obtain:

$\displaystyle\log p_{\bm{\theta}}(\mathbf{x})$
$\displaystyle\geq\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{y})]-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{y}))$
(1)
$\displaystyle\geq\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{z}|\mathbf{x},\mathbf{y})||p_{\bm{\theta}}(\mathbf{z}|\mathbf{y}))]-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{y}))$
(2)

While Equation 2 is a valid learning objective, it does not yield a training
procedure comparable to that of a diffusion model. Diffusion model training
involves sampling noisy data from the forward process; here, we sample from
the approximate reverse process (from $q_{\bm{\phi}}(\mathbf{z}|\mathbf{y})$)
and seek to match the true process $r(\mathbf{z}|\mathbf{y})$. In practice, we
have found the learning signal from this procedure to be too weak to learn a
good $q_{\bm{\phi}}(\mathbf{z}|\mathbf{y})$ that reverses the noising process
$r(\mathbf{y}|\mathbf{z})$, as illustrated in Table 7.

Instead, we adopt a learning objective
$\mathcal{L}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$ inspired by the wake-sleep
algorithm:

$\displaystyle\mathcal{L}$
$\displaystyle=\underbrace{\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})]}_{\text{wake / recons. term
$\mathcal{L}_{\text{rec}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$}}-\underbrace{D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{y},\mathbf{z}))}_{\text{prior
regularization term
$\mathcal{L}_{\text{reg}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$}}-\underbrace{\mathbb{E}_{p_{\bm{\theta}}(\mathbf{x})}[D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))]}_{\text{sleep
term $\mathcal{L}_{\text{sleep}}({\bm{\phi}})$}}$ (3)

Observe that this objective is the ELBO in Equation 2 augmented with an
additional regularizer $\mathcal{L}_{\text{sleep}}({\bm{\phi}})$ that is
inspired by the sleep phase of the wake-sleep algorithm. This term consists of
the reverse KL divergence
$D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))$.
Section 3.1.2 shows that this divergence can be optimized via a diffusion-like
training procedure. As in wake-sleep, we optimize $\mathcal{L}_{\text{sleep}}$
over ${\bm{\phi}}$ only.

Objective (3) poses restrictions on $p_{\bm{\theta}}$ and $r$. The prior
$p_{\bm{\theta}}(\mathbf{y},\mathbf{z})$ needs to have a tractable density,
although we will also present approximations and experimental results with
implicit sample-based priors. Additionally,
$q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})$ must feature tractable
entropy and sampling.

Lastly, note that $\mathcal{L}$ lower bounds the marginal log-likelihood $\log
p_{\bm{\theta}}(\mathbf{x})$. When
$q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})$ equals the true posterior,
this bound is tight, since the ELBO is tight, and the sleep term also equals
zero.

####  3.1.2 Optimization Using Wake-Sleep in Latent Space

Next, we introduce optimization algorithms for diffusion variational
inference. Maximizing $\mathcal{L}(x,{\bm{\theta}},{\bm{\phi}})$ involves
optimizing the sleep term $\mathcal{L}_{\text{sleep}}({\bm{\phi}})$. This
optimization is tractable since:

$\displaystyle\mathcal{L}_{\text{sleep}}({\bm{\phi}})$
$\displaystyle=-\mathbb{E}_{p_{\bm{\theta}}(\mathbf{x})}[D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})|q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))]=\mathbb{E}_{p_{\bm{\theta}}(\mathbf{x},\mathbf{z})}[\log[{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}/{p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})}]]$
(4)
$\displaystyle=\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})}[\log{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}]+\bar{H}(p_{\bm{\theta}})$
(5)
$\displaystyle\geq\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})}[\mathbb{E}_{r(\mathbf{y}|\mathbf{z})}[\log(q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})/r(\mathbf{y}|\mathbf{z}))]]+\bar{H}(p_{\bm{\theta}})$
(6)

where in Equation 6 we applied the ELBO with $r(\mathbf{y}|\mathbf{z})$ as our
choice of variational posterior over the latent $\mathbf{y}$ in the
distribution $q_{\bm{\phi}}$; $\bar{H}(p_{\bm{\theta}})$ is the expected
conditional entropy of $p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})$, a constant
that does not depend on ${\bm{\phi}}$. By maximizing Equation 6, we fit
$q(\mathbf{z}|\mathbf{x})$ to $p_{\bm{\theta}}(\mathbf{z}|\mathbf{x})$ while
also fitting $q(\mathbf{z}|\mathbf{y})$ to the true reverse denoising process
$r(\mathbf{z}|\mathbf{y})$ via the reconstruction term
$\mathbb{E}_{r(\mathbf{y}|\mathbf{z})}[\log
q_{\bm{\phi}}(\mathbf{z}|\mathbf{x},\mathbf{y})]$, which is one of the terms
comprising Equation 6. To fit the objective
$\mathcal{L}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$, we propose introducing an
additional sleep step in which we sample from
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})$ in order to compute
gradients for $\mathcal{L}_{\text{sleep}}({\bm{\phi}})$.

This procedure mirrors wake-sleep; however, sampling $\mathbf{x}$ from
$p_{\bm{\theta}}$ to obtain gradients for the sleep term introduces
computational overhead. To address this issue, we propose wake-sleep in latent
space, an algorithm that optimizes an approximation
$\hat{\mathcal{L}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$ of $\mathcal{L}$:

$\displaystyle\hat{\mathcal{L}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$
$\displaystyle=\underbrace{\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})]}_{\text{wake / reconstr. term
$\mathcal{L}_{\text{rec}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$}}-\underbrace{D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{y},\mathbf{z}))}_{\text{prior
regularization term
$\mathcal{L}_{\text{reg}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$}}-\underbrace{D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z})||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))}_{\text{latent
sleep term $\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}})$}}.$ (7)

We have replaced $\mathcal{L}_{\text{sleep}}({\bm{\phi}})$ with a latent sleep
term $\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}})$, in which
$\mathbf{x}$ is given, and we only seek to fit the true reverse noising
process $r(\mathbf{z}|\mathbf{y})$ independently of $\mathbf{x}$. We can
similarly show that

$\displaystyle\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}})$
$\displaystyle=\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})}[\log{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}]+\bar{H}(p_{\bm{\theta}})\geq\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})r(\mathbf{y}|\mathbf{z})}[\log({q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})}/r(\mathbf{y}|\mathbf{z}))]+\bar{H}(p_{\bm{\theta}})$
(8)
$\displaystyle=-\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})}[D_{\mathrm{KL}}(r(\mathbf{y}|\mathbf{z})||q_{\bm{\phi}}(\mathbf{y}|\mathbf{z},\mathbf{x}))]-D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z})||q(\mathbf{z}|\mathbf{x})),$
(9)

where $\bar{H}(p_{\bm{\theta}})$ is an entropy term constant in ${\bm{\phi}}$.
Thus, we minimize the forward KL divergence by sampling $\mathbf{z}$, and
applying the noising process to get $\mathbf{y}$; the $q_{\bm{\phi}}$ is fit
to denoise $\mathbf{z}$ from $\mathbf{y}$ as in Equation 6.

We optimize our bound on $\hat{\mathcal{L}}(x,{\bm{\theta}},{\bm{\phi}})$ end-
to-end using minibatch gradient descent over ${\bm{\theta}},{\bm{\phi}}$.
While the wake term is a reconstruction loss as in wake-sleep, the sleep term
generates latent samples $\mathbf{z},\mathbf{y}$ from
$r(\mathbf{y}|\mathbf{z})p_{\bm{\theta}}(\mathbf{z})$ (by analogy with
$p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})p_{\bm{\theta}}(\mathbf{z})$ in normal
wake-sleep); the denoiser $q_{\bm{\phi}}$ is trained to recover $\mathbf{z}$
from $\mathbf{y}$. Thus, we perform wake-sleep in latent space, which obviates
the need for alternating wake and sleep phases, and allows efficient end-to-
end training. A limitation of this approximation is that the sleep term does
not fit $q_{\bm{\phi}}$ to the true
$p_{\bm{\theta}}(\mathbf{z}|\mathbf{x},\mathbf{y})$, and as a consequence
$\hat{L}$ is not a tight lower bound on $\log p_{\bm{\theta}}(\mathbf{x})$. We
may think of $\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}})$ as a
regularizer to the ELBO.

####  3.1.3 DDVI with Diffusion-Based Encoders

Our framework is naturally instantiated with diffusion models. Let
$\mathbf{y}_{0}=\mathbf{z}$ and $\mathbf{y}_{T}=\mathbf{y}$. The forward
noising model can be defined as
$r(\mathbf{y}_{1:T}|\mathbf{z})=\prod_{t=1}^{T}r(\mathbf{y}_{t}|\mathbf{y}_{t-1})$,
where the $\mathbf{y}_{0:T}$ are an extended set of latent variables that
represent increasingly noised versions of $\mathbf{y}_{0}$. We parameterize
the approximate reverse diffusion process as
$q_{\bm{\phi}}(\mathbf{y}_{0:T}|\mathbf{x})=q(\mathbf{y}_{T}|\mathbf{x})\prod_{t=1}^{T}q_{\bm{\phi}}(\mathbf{y}_{t-1}|\mathbf{y}_{t},\mathbf{x}).$

We can form a lower bound $\mathcal{L}_{\text{diff}}(\mathbf{x},{\bm{\phi}})$
on the sleep term
$\mathbb{E}_{p(\mathbf{z})}\log{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}$ in
Equation 8 (where $q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})=\int
q_{\bm{\phi}}(\mathbf{z},\mathbf{y}_{1:T}|\mathbf{x})d\mathbf{y}_{1:T}$) using
the ELBO for a diffusion model:

$\displaystyle\mathcal{L}_{\text{diff}}=$
$\displaystyle\mathbb{E}_{r}\left[\log
q_{\bm{\phi}}(\mathbf{z}|\mathbf{y}_{1},\mathbf{x})-\sum_{t=2}^{T}D_{\mathrm{KL}}(r(\mathbf{y}_{t-1}|\mathbf{y}_{t},\mathbf{y}_{0})||q_{\bm{\phi}}(\mathbf{y}_{t-1}|\mathbf{y}_{t},\mathbf{x}))\right]-D_{\mathrm{KL}}(r(\mathbf{y}_{T}|\mathbf{z})||q_{\bm{\phi}}(\mathbf{y}|\mathbf{x})).$
(10)

This bound is an instantiation of Equation 8 when $r$ is a diffusion process.
Many choices of $r,q$ fit this framework. A common type of noising process is
Gaussian diffusion, where we define
$r(\mathbf{y}_{t}|\mathbf{y}_{t-1})=\mathcal{N}(\mathbf{y}_{t};\sqrt{1-{\bm{\alpha}}_{t}}\mathbf{y}_{t-1},{\bm{\alpha}}_{t}\mathbf{I})$
for a suitable schedule $({\bm{\alpha}}_{t})_{t=1}^{T}$. We then adopt the
parameterization
$q_{\bm{\phi}}(\mathbf{y}_{t-1}|\mathbf{y}_{t},\mathbf{x})=\mathcal{N}(\mathbf{y}_{t-1};\mu_{\bm{\phi}}(\mathbf{y}_{t},\mathbf{x},t),\Sigma_{\bm{\phi}}(\mathbf{y}_{t},\mathbf{x},t))$.
It is then common to parameterize $q_{\bm{\phi}}$ with a noise prediction
network ${\bm{\epsilon}}_{\bm{\phi}}$ (Ho et al., 2020); the sum of KL
divergences can be approximated by $E_{t,{\bm{\epsilon}}_{t}\sim
r(\mathbf{y}_{0},t)}||{\bm{\epsilon}}_{t}-{\bm{\epsilon}}_{\bm{\phi}}(\sqrt{\bar{{\bm{\alpha}}_{t}}}\mathbf{y}_{0}+\sqrt{1-\bar{{\bm{\alpha}}_{t}}}{\bm{\epsilon}}_{t},\mathbf{x},t)||^{2}$.
By our earlier argument, this objective encourages
$q_{\bm{\phi}}(\mathbf{y}_{0:T}|\mathbf{x})$ to match $r(\mathbf{y}_{0:T})$ as
well as the true posterior $p(\mathbf{z}|\mathbf{x})$.

Lastly, we need to show that this choice of $q$ admits a tractable entropy.
This follows from

$\displaystyle H(q)$ $\displaystyle=-\sum_{t=1}^{T+1}\mathbb{E}_{q}[\log
q(\mathbf{y}_{t-1}|\mathbf{y}_{t},\mathbf{x})]=\sum_{t=1}^{T+1}\mathbb{E}_{q}[\frac{d}{2}\left(1+\log(2\pi)\right)+\frac{1}{2}\log|\Sigma_{\bm{\phi}}(\mathbf{y}_{t},\mathbf{x})|]$
(11)

where $d$ is the dimension of $\mathbf{y}$ and we use the notation
$\mathbf{y}_{T+1}=\mathbf{x}$. The right-hand term can be approximated using
Monte Carlo; it is also common to leave the variance $\Sigma_{\bm{\phi}}$
fixed (as we typically do in our experiments), in which case $H(q)$ is a
constant.

###  3.2 Auxiliary-Variable Diffusion Encoders

Denoising diffusion variational inference can also be understood from the
perspective of auxiliary-variable generative models (Maaløe et al., 2016). Our
method can be seen as introducing an approximate posterior
$q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})=\int_{\mathbf{y}}q_{\bm{\phi}}(\mathbf{z},\mathbf{y}|\mathbf{x})$
which itself contains latent variables $\mathbf{y}$.

The standard approach to fit auxiliary-variable generative models is to apply
the ELBO twice:

$\displaystyle\log p_{\bm{\theta}}(\mathbf{x})$ $\displaystyle\geq\log
p_{\bm{\theta}}(\mathbf{x})-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{z}|\mathbf{x}))=\text{ELBO}_{\mathbf{z}|\mathbf{x}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})$
(12)
$\displaystyle\geq\text{ELBO}_{\mathbf{z}|\mathbf{x}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})-\mathbb{E}_{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}[D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y}|\mathbf{x},\mathbf{z})||r(\mathbf{y}|\mathbf{x},\mathbf{z}))]$
(13)
$\displaystyle=\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x},\mathbf{z})r(\mathbf{y}|\mathbf{x},\mathbf{z})-\log
q(\mathbf{y},\mathbf{z}|\mathbf{x})],$ (14)

where in Equation 13 we applied the ELBO over $\mathbf{z}$, and in Equation 12
we applied the ELBO again over the latent $\mathbf{y}$ of $q$. The
$r(\mathbf{y}|\mathbf{x},\mathbf{z})$ can be interepreted as an approximate
variational posterior for the true posterior
$q_{\bm{\phi}}(\mathbf{y}|\mathbf{x},\mathbf{z})$ of the data distribution.
Specifically, the gap between Equation 12 and $\log
p_{\bm{\theta}}(\mathbf{x})$ is precisely
$D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{z}|\mathbf{x}))+\mathbb{E}_{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})}[D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y}|\mathbf{x},\mathbf{z})||r(\mathbf{y}|\mathbf{x},\mathbf{z}))]$.
Thus, if we correctly match the pairs of approximate posteriors, we will
achieve a tight bound.

Additionally, notice that Equation 14 is exactly the ELBO in Equation 2 after
rearranging terms. Thus the two perspectives are equivalent: using variational
inference to fit an augmented model
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z})$ is equivalent to fitting
the original $p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$ and applying the ELBO to
$p_{\bm{\theta}}$, followed by applying the ELBO again to the latent
$\mathbf{y}$ in $q$. Since our objective $\mathcal{L}$ is Equation 14)
augmented with $\mathcal{L}_{\text{sleep}}$, denoising VI can also be
interpreted as fitting an latent-variable approximate posterior $q$. As seen
in Equation 9, we implicitly minimize the reverse KL divergence
$\mathbb{E}_{p_{\bm{\theta}}(\mathbf{z})}[D_{\mathrm{KL}}(r(\mathbf{y}|\mathbf{z})||q_{\bm{\phi}}(\mathbf{y}|\mathbf{z},\mathbf{x}))]$.

Note, however, that there is a crucial difference between our method and
auxiliary-variable generative models: we do not optimize the parameters of
$r$. Rather, $r$ is a user-defined noising process that acts as a regularizer
for the model $q_{\bm{\phi}}$ and whose reverse we seek to fit.

###  3.3 Regularizing Variational Inference with Denoising Processes

Lastly, our approach can be understood from a third perspective: that of
regularization. As we mentioned earlier, $\mathcal{L}$ is the ELBO with an
additional regularization term that encourages the approximate posterior
$q(\mathbf{z}|\mathbf{x},\mathbf{y})q(\mathbf{y}|\mathbf{x})$ to gradually
transform $\mathbf{y}$ into $\mathbf{z}$ following the guidance of $r$. In
that sense, $r$ provides data augmentation to ensure that the mapping from
$\mathbf{y}$ to $\mathbf{x}$ is easier to learn, similarly to how a diffusion
process provides additional learning signal for generating $\mathbf{x}$ from
noise.

Our framework also supports the user defining general distributions as
regularizers $r_{1}(\mathbf{z},\mathbf{y})$ and
$r_{2}(\mathbf{y}|\mathbf{z},\mathbf{x})$. Here, $r_{1}$ puts a custom prior
on $\mathbf{z},\mathbf{y}$, while $r_{2}$ specifies a mapping between
$\mathbf{y},\mathbf{z}$ that should be followed by the encoder model. We write
$\mathcal{L}_{\text{reg}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}},r_{1})$ and
$\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}},r_{2})$, where each term is
still defined as in Equation 3 and Equation 6, except using $r_{1},r_{2}$ and
we optimize

$\displaystyle\mathcal{L}(\mathbf{x},{\bm{\theta}},{\bm{\phi}},r_{1},r_{2})=\mathcal{L}_{\text{rec}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}})+\beta\cdot\mathcal{L}_{\text{reg}}(\mathbf{x},{\bm{\theta}},{\bm{\phi}},r_{1})+\gamma\cdot\mathcal{L}_{\text{sleep}}(\mathbf{x},{\bm{\phi}},r_{2}),$
(15)

where $\beta,\gamma>0$ (in particular, $\beta$ mirrors the $\beta$-VAE model
(Higgins et al., 2016)).

Examples of $r_{2}$ include Gaussian diffusion, discrete noising processes, as
well as noise conditioned on $\mathbf{x}$. Note that $r_{2}$ could even be
fully deterministic when parameterized with a bijective flow. With a strong
regularizer, we could set $q(\mathbf{z}|\mathbf{y},\mathbf{x})$ to equal the
reverse $r_{2}(\mathbf{z}|\mathbf{y},\mathbf{x})$ of this flow. Thus, for
suitably designed $r_{2}$, we could perfectly match the shape of any mapping.
This mirrors the InteL-VAE framework (Miao et al., 2021), with support for
flexible and probabilistic regularizers.

Most importantly, our method can be understood as introducing a new form of
regularization. While a standard VAE introduces a prior regularization term,
we introduce a new form of regularization via a function $r_{2}$. This can be
advantageous when specifying or enforcing certain priors is difficult.

##  4 Extensions

###  4.1 Semi-Supervised Learning

Following Makhzani et al. (2015), we extend our algorithm to the semi-
supervised learning setting where some data points are labeled and we define
$p(\mathbf{z},l)=p(\mathbf{z}|c)p(l)$. Then, the model can be specified as
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z},l)=p_{\bm{\theta}}(\mathbf{x}|\mathbf{z},l)r(\mathbf{y}|\mathbf{z},l)p_{\bm{\theta}}(\mathbf{z}|l)p(l)$
and the variational distributions are
$q_{\bm{\phi}}(\mathbf{z}|\mathbf{x},\mathbf{y},l),q_{\bm{\phi}}(\mathbf{y}|\mathbf{x}),q_{\bm{\phi}}(l|\mathbf{x})$.
In this setting, we consider two cases of whether the label can be observed or
not (Kingma et al., 2014). We extend Equation 7 to incorporate the label $l$
corresponding to a data point as follows:

$\displaystyle\mathcal{L}_{\text{semi}}=\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x},l)}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z},l)]-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y},\mathbf{z}|\mathbf{x},l)||p_{\bm{\theta}}(\mathbf{y},\mathbf{z}|l))-D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z}|l)||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x},l))$

When the label $c$ cannot be observed, we treat it as a latent variable and
modify the learning objective
$\mathcal{U}_{\text{semi}}=\sum_{c}q_{\bm{\phi}}(l|\mathbf{x})\mathcal{L}_{\text{semi}}(\mathbf{x},l,{\bm{{\bm{\theta}}}},{\bm{\phi}})+D_{\mathrm{KL}}(q_{\bm{\phi}}(l|\mathbf{x})||p(l))$.
Therefore, we can conclude a marginal likelihood on our dataset as follows:
$\tilde{\mathcal{L}}_{\text{semi}}=\sum_{(\mathbf{x},l)\in
L}\mathcal{L}_{\text{semi}}(\mathbf{x},l,{\bm{{\bm{\theta}}}},{\bm{{\bm{\phi}}}})+\sum_{\mathbf{x}\in
U}\mathcal{U}_{\text{semi}}(\mathbf{x},{\bm{{\bm{\theta}}}},{\bm{{\bm{\phi}}}}).$
where $L$ and $U$ are the sets of data with and without labels, respectively.

We also want to guarantee that all model parameters can be learned in all
cases, including $q_{\bm{\phi}}(l|\mathbf{x})$, such that this posterior can
be applied as a classifier during inference. Thus, we combine the marginal
likelihood with a classification loss to form an extended learning objective:
$\tilde{\mathcal{L}}_{\text{semi}_{\bm{\alpha}}}=\tilde{\mathcal{L}}_{\text{semi}}+{\bm{\alpha}}\cdot\mathbb{E}_{\tilde{p}(\mathbf{x},l)}\left[-\log
q_{\bm{\phi}}(l|\mathbf{x})\right]$

###  4.2 Clustering

We have further extended our algorithm to encompass the clustering paradigm.
We propose two distinct strategies. In the first approach, we simply set
$p_{\bm{{\bm{\theta}}}}(\mathbf{z})$ as a mixture of desired priors. The means
of these priors are characterized by ${\bm{{\bm{\theta}}}}$. From these means,
cluster membership, denoted as $\mathbf{c}$ can be deduced. Remarkably, this
approach requires no alteration to the existing learning objective.

Alternatively, the second method retains the original prior but introduces an
additional cluster latent variable $\mathbf{c}$ where $\sum_{i}c_{i}=1$. Thus,
the model can be specified as
$p_{\bm{\theta}}(\mathbf{x},\mathbf{y},\mathbf{z},\mathbf{c})=p_{\bm{\theta}}(\mathbf{x}|\mathbf{z},\mathbf{c})r(\mathbf{y}|\mathbf{z})p_{\bm{\theta}}(\mathbf{z})p(\mathbf{c})$
with $p(\mathbf{c})=Dir({\bm{{\bm{\epsilon}}}})$. Consequently, the
variational distributions become
$q_{\bm{\phi}}(\mathbf{z}|\mathbf{y},\mathbf{c},\mathbf{x}),q_{\bm{\phi}}(\mathbf{y},\mathbf{c}|\mathbf{x})$.
This reformulates the learning objective as:

$\displaystyle\mathcal{L}_{\text{clus}}(\mathbf{x})$
$\displaystyle=\mathbb{E}_{q_{\bm{\phi}}(\mathbf{y},\mathbf{z},\mathbf{c}|\mathbf{x})}[\log
p_{\bm{\theta}}(\mathbf{x}|\mathbf{z},\mathbf{c})]-D_{\mathrm{KL}}(q_{\bm{\phi}}(\mathbf{y},\mathbf{z},\mathbf{c}|\mathbf{x})||p_{\bm{\theta}}(\mathbf{y},\mathbf{z},\mathbf{c}))-D_{\mathrm{KL}}(p_{\bm{\theta}}(\mathbf{z})||q_{\bm{\phi}}(\mathbf{z}|\mathbf{x}))$

##  5 Experiments

We compare DD-VAE with vanilla VAE (Kingma & Welling, 2013), IAF-VAE (Kingma
et al., 2016), and AAE (Makhzani et al., 2015) on MNIST (Lecun et al., 1998)
and CIFAR-10 (Krizhevsky & Hinton, 2009) in unsupervised and semi-supervised
learning settings, and also on the 1000 genomes dataset Siva (2008) in
clustering settings. The differences between these methods are summarized in
Table 8. We discuss the computational costs of all methods in Section A.1. The
priors, model architecture, and training details can be founded in Section
A.2, Section A.3, and Section A.4 respectively. All results below are reported
with 95% confidence interval using 5 different seeds.

###  5.1 Unsupervised learning

We fit a model $p_{\bm{\theta}}(\mathbf{x},\mathbf{z})$ on the MNIST and
CIFAR-10 datasets with three priors $p(\mathbf{z})$: pinwheel, swiss roll, and
square and report our results in Table 1 and Table 5. We measure a k-nearest
neighbors classification accuracy of the latents (Acc). We also measure latent
negative log-likelihood (Latent NLL) by fitting a kernel density estimator
(KDE) on the latents produced by the model with test data as input and compute
the log-likelihood of the latents sampled from the prior under the fitted KDE.

We consistently see our method perform best on both KNN accuracy and latent
NLL, indicating that its output latents are the most informative of labels and
best-aligned with the priors. Table 1 reports qualiatitve results that
indicate that DD-VAE enforces most accurately the correct prior shape.

Figure 2: Unsupervised visualization on MNIST using three different priors
(pinwhell, swiss roll, and square). Each color indicates a class.

Method Pinwheel Swiss Roll Square Acc Latent NLL Acc Latent NLL Acc Latent NLL
VAE $0.58\pm 0.01$ $6.18\pm 0.93$ $0.56\pm 0.055$ $82.96\pm 99.30$ $0.35\pm
0.02$ $25.07\pm 62.41$ IAF-VAE $\bm{0.71\pm 0.00}$ $2.06\pm 0.21$ $0.67\pm
0.02$ $12.13\pm 18.28$ $0.65\pm 0.03$ $1.53\pm 0.82$ AAE $0.67\pm 0.02$
$1.94\pm 0.09$ $0.68\pm 0.03$ $3.43\pm 0.14$ $0.57\pm 0.02$ $1.97\pm 0.70$
HVAE w/ flow $\bm{0.72\pm 0.01}$ $2.18\pm 0.13$ $\bm{0.71\pm 0.02}$ $4.17\pm
1.23$ $\bm{0.70\pm 0.01}$ $\bm{1.25\pm 0.08}$ DiffVAE $\bm{0.72\pm 0.01}$
$\bm{1.57\pm 0.10}$ $0.66\pm 0.01$ $\bm{3.19\pm 0.07}$ $\bm{0.72\pm 0.03}$
$\bm{1.24\pm 0.13}$

Table 1: Unsupervised learning on MNIST. We report accuracy using KNN (K=20)
classifier and latent negative log-likelihood with pinwheel, swiss roll, and
square priors.

###  5.2 Semi-supervised Learning

We also evaluate the performance of our method and the baselines under semi-
supervised learning setting where some labels are observed (1,000 for MNIST
and 10,000 for CIFAR-10) and the partitions of the priors are known. We use
the same set of priors and baselines. Details on how we partition each prior
into $p_{\bm{{\bm{\theta}}}}(\mathbf{z}|l)$ can be founded in Section A.2. The
partitions defined for our priors are local parts of the priors. We substitute
the KNN accuracy metric with the average class latent NLL, which measures how
well the test latents from each class match their respective partition of the
prior $p_{\bm{{\bm{\theta}}}}(\mathbf{z}|l)$. In our case, this metric
implicitly measures accuracy because the metric requires that the test latents
for each class be close to each other, as the partitions defined are local
parts of the full prior. The results are shown in Table 2 and Table 6. DD-VAE
mostly outperforms the baselines across different priors and metrics. We also
show the visualization in Figure 3 where DD-VAE matches the prior almost
perfectly.

Figure 3: Semi-supervised visualization on MNIST with 1,000 labels using three
different priors (pinwheel, swiss roll, and square). Each a indicates one
class.

Method Pinwheel Swiss Roll Square Latent NLL Avg Class Latent NLL Avg Class
Latent NLL Avg Class Latent NLL Latent NLL Latent NLL VAE $7.77\pm 1.83$
$15.64\pm 8.79$ $56.91\pm 33.90$ $244.51\pm 94.19$ $1.30\pm 0.03$ $-0.71\pm
0.05$ IAF-VAE $2.71\pm 0.22$ $1.84\pm 0.46$ $8.10\pm 12.17$ $14.59\pm 31.88$
$1.04\pm 0.04$ $-1.01\pm 0.03$ AAE $1.90\pm 0.04$ $0.64\pm 0.11$ $\bm{3.35\pm
0.17}$ $1.86\pm 0.77$ $1.70\pm 0.26$ $2.50\pm 4.70$ DiffVAE $\bm{1.28\pm
0.01}$ $\bm{-0.94\pm 0.01}$ $\bm{3.35\pm 0.10}$ $\bm{1.24\pm 0.09}$
$\bm{0.96\pm 0.03}$ $\bm{-1.20\pm 0.02}$

Table 2: Semi-supervised learning on MNIST (1,000 labels). We report average
class latent negative log-likelihood and latent negative log-likelihood with
pinwheel, swiss roll, and square priors.

###  5.3 Clustering and Visualization for Genotype Analysis

Figure 4: Visualization of genotype clusters. A color represents one
ethnicity.

In this section, we report results on an real-world task in genome analysis.
Visualizing genotype data reveals patterns in the latent ancestry of
individuals. We compare DD-VAE against with the three strong clustering
baselines using the 1000 Genomes dataset. We also report visualizations from
three dimensionality reduction algorithms: PCA, TSNE, and UMAP. For each
clustering algorithm, we seek to discover up to 20 clusters. We report
quantitative results in terms of cluster purity, cluster completeness, and
normalized mutual information (NMI). There is an inherent trade-off between
cluster purity completeness. The overall clustering performance can be
captured with NMI.

Method | Cluster | Cluster | NMI
---|---|---|---
Method | Purity | Completeness |
VAE | $0.28\pm 0.02$ | $\bm{0.78\pm 0.16}$ | $0.59\pm 0.08$
IAF-VAE | $0.29\pm 0.04$ | $\bm{0.73\pm 0.06}$ | $0.55\pm 0.06$
AAE | $0.37\pm 0.06$ | $\bm{0.76\pm 0.11}$ | $\bm{0.63\pm 0.02}$
DiffVAE | $\bm{0.45\pm 0.03}$ | $\bm{0.75\pm 0.05}$ | $\bm{0.66\pm 0.04}$

Table 3: Quantitative genotype clustering results.

In Table 3, we see that DD-VAE attains the best performance on cluster purity
and NMI. For cluster completeness, VAE and AAE have better means but much
larger confidence interval. Furthermore, we visualize our genotype clustering
results in latent space, shown in Figure 4, and also report results from
classical dimensionality reduction and visualization methods that do not
perform clustering (PCA (Wold et al., 1987), t-SNE (Van der Maaten & Hinton,
2008), and UMAP (McInnes et al., 2018)). The legend of Figure 4 can be founded
at Figure 5.

##  6 Related Work, Discussion, and Conclusion

There are a number of recent works merging VAE and diffusion models. Vahdat et
al. (2021); Wehenkel & Louppe (2021) use diffusion with VAE priors. Recent
research (Preechakul et al., 2022; Zhang et al., 2022; Wang et al., 2023b) has
also melded auto-encoders with diffusion models, focusing on semantically
meaningful latents. Cohen et al. (2022) crafts a diffusion bridge linking a
continuous coded vector to a non-informative prior distribution. These works
focus on sample quality and directly fit
$q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})$ to a complex distribution (e.g.,
compressed images (Rombach et al., 2022)); our work instead seeks to achieve
high sample quality in latent space, and obtains complex $\mathbf{z}$ by
diffusion from a simple $\mathbf{y}\sim q_{\bm{\phi}}(\mathbf{y}|\mathbf{x})$.

Latent variable models in general are an attractive alternative to
visualization methods like PCA, UMAP, and t-SNE (McInnes et al., 2018; Van der
Maaten & Hinton, 2008). Domain-specific knowledge can be injected through the
prior, and deep neural networks can be utilized to achieve a more expressive
mapping from the data space to the latent space. Nevertheless, downsides of
LVMs are that they are more computationally expensive and require careful
hyperparameter tuning.

While this paper focuses on applications of DD-VAE to dimensionality reduction
and visualization, there exist other tasks for the algorithm, e.g., density
estimation or sample quality. Since our learning objective differs from the
ELBO (it adds a regularizer), we anticipate gains on models whose training
benefits from regularization, but perhaps not on all models. Also, attaining
competitive likelihood estimation requires architecture improvements that are
orthogonal to this paper. However, our ability to generate diverse samples and
achieve class separation in latent space hints at the method’s potential on
these tasks. Accurate variational inference has the potential to improve
downstream applications of generative modeling, e.g., decision making or
causal effect estimation (Nguyen & Grover, 2022; Deshpande et al., 2022;
Deshpande & Kuleshov, 2023; Rastogi et al., 2023).

## References

  * Amid & Warmuth (2019) Ehsan Amid and Manfred K Warmuth.  Trimap: Large-scale dimensionality reduction using triplets.  _arXiv preprint arXiv:1910.00204_ , 2019.
  * Austin et al. (2021) Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg.  Structured denoising diffusion models in discrete state-spaces.  _Advances in Neural Information Processing Systems_ , 34:17981–17993, 2021.
  * Balakrishnama & Ganapathiraju (1998) Suresh Balakrishnama and Aravind Ganapathiraju.  Linear discriminant analysis-a brief tutorial.  _Institute for Signal and information Processing_ , 18(1998):1–8, 1998.
  * Balasubramanian & Schwartz (2002) Mukund Balasubramanian and Eric L Schwartz.  The isomap algorithm and topological stability.  _Science_ , 295(5552):7–7, 2002.
  * Battey et al. (2021) CJ Battey, Gabrielle C Coffing, and Andrew D Kern.  Visualizing population structure with variational autoencoders.  _G3_ , 11(1):jkaa036, 2021.
  * Belkin & Niyogi (2003) Mikhail Belkin and Partha Niyogi.  Laplacian eigenmaps for dimensionality reduction and data representation.  _Neural computation_ , 15(6):1373–1396, 2003.
  * Carroll & Arabie (1998) J Douglas Carroll and Phipps Arabie.  Multidimensional scaling.  _Measurement, judgment and decision making_ , pp. 179–250, 1998.
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers, 2019.
  * Cohen et al. (2022) Max Cohen, Guillaume Quispe, Sylvain Le Corff, Charles Ollion, and Eric Moulines.  Diffusion bridges vector quantized variational autoencoders.  _arXiv preprint arXiv:2202.04895_ , 2022.
  * Coifman & Lafon (2006) Ronald R Coifman and Stéphane Lafon.  Diffusion maps.  _Applied and computational harmonic analysis_ , 21(1):5–30, 2006.
  * Deshpande & Kuleshov (2023) Shachi Deshpande and Volodymyr Kuleshov.  Calibrated uncertainty estimation improves bayesian optimization, 2023.
  * Deshpande et al. (2022) Shachi Deshpande, Kaiwen Wang, Dhruv Sreenivas, Zheng Li, and Volodymyr Kuleshov.  Deep multi-modal structural equations for causal effect estimation with unstructured proxies.  _Advances in Neural Information Processing Systems_ , 35:10931–10944, 2022.
  * Diaz-Papkovich et al. (2019) Alex Diaz-Papkovich, Luke Anderson-Trocmé, Chief Ben-Eghan, and Simon Gravel.  Umap reveals cryptic population structure and phenotype heterogeneity in large genomic cohorts.  _PLoS genetics_ , 15(11):e1008432, 2019.
  * Goodfellow et al. (2014) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  _Advances in neural information processing systems_ , 27, 2014.
  * Haghverdi et al. (2015) Laleh Haghverdi, Florian Buettner, and Fabian J Theis.  Diffusion maps for high-dimensional single-cell analysis of differentiation data.  _Bioinformatics_ , 31(18):2989–2998, 2015.
  * Henderson (1997) Paul Henderson.  Sammon mapping.  _Pattern Recognit. Lett_ , 18(11-13):1307–1316, 1997.
  * Higgins et al. (2016) Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  beta-vae: Learning basic visual concepts with a constrained variational framework.  In _International conference on learning representations_ , 2016.
  * Hinton et al. (1995) Geoffrey E Hinton, Peter Dayan, Brendan J Frey, and Radford M Neal.  The” wake-sleep” algorithm for unsupervised neural networks.  _Science_ , 268(5214):1158–1161, 1995.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Im Im et al. (2017) Daniel Im Im, Sungjin Ahn, Roland Memisevic, and Yoshua Bengio.  Denoising criterion for variational auto-encoding framework.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 31, 2017.
  * Johnson et al. (2016) Matthew J. Johnson, David Duvenaud, Alexander B. Wiltschko, Sandeep Robert Datta, and Ryan P. Adams.  Composing graphical models with neural networks for structured representations and fast inference.  In _Advances in Neural Information Processing Systems (NIPS) 29_ , 2016.  arXiv:1603.06277 [stat.ML].
  * Kingma et al. (2021) Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  _Advances in neural information processing systems_ , 34:21696–21707, 2021.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kingma et al. (2014) Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling.  Semi-supervised learning with deep generative models.  _Advances in neural information processing systems_ , 27, 2014.
  * Kingma et al. (2016) Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  _Advances in neural information processing systems_ , 29, 2016.
  * Krizhevsky & Hinton (2009) Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  2009\.
  * Lecun et al. (1998) Yann Lecun, Leon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.  doi: 10.1109/5.726791.
  * Maaløe et al. (2016) Lars Maaløe, Casper Kaae Sønderby, Søren Kaae Sønderby, and Ole Winther.  Auxiliary deep generative models.  In _International conference on machine learning_ , pp. 1445–1453. PMLR, 2016.
  * Makhzani et al. (2015) Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey.  Adversarial autoencoders.  _arXiv preprint arXiv:1511.05644_ , 2015.
  * Marsland (2014) Stephen Marsland.  _Machine Learning: An Algorithmic Perspective (2nd Edition)_.  Chapman and Hall/CRC, 2014.
  * McInnes et al. (2018) Leland McInnes, John Healy, and James Melville.  Umap: Uniform manifold approximation and projection for dimension reduction.  _arXiv preprint arXiv:1802.03426_ , 2018.
  * Miao et al. (2021) Ning Miao, Emile Mathieu, N Siddharth, Yee Whye Teh, and Tom Rainforth.  On incorporating inductive biases into vaes.  _arXiv preprint arXiv:2106.13746_ , 2021.
  * Moon et al. (2017) Kevin R Moon, David van Dijk, Zheng Wang, William Chen, Matthew J Hirn, Ronald R Coifman, Natalia B Ivanova, Guy Wolf, and Smita Krishnaswamy.  Phate: a dimensionality reduction method for visualizing trajectory structures in high-dimensional biological data.  _BioRxiv_ , 120378, 2017.
  * Nguyen & Grover (2022) Tung Nguyen and Aditya Grover.  Transformer neural processes: Uncertainty-aware meta learning via sequence modeling.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), _International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 16569–16594. PMLR, 2022.  URL https://proceedings.mlr.press/v162/nguyen22b.html.
  * Nichol & Dhariwal (2021) Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In _International Conference on Machine Learning_ , pp. 8162–8171. PMLR, 2021.
  * Preechakul et al. (2022) Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn.  Diffusion autoencoders: Toward a meaningful and decodable representation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 10619–10629, 2022.
  * Rastogi et al. (2023) Richa Rastogi, Yair Schiff, Alon Hacohen, Zhaozhi Li, Ian Lee, Yuntian Deng, Mert R Sabuncu, and Volodymyr Kuleshov.  Semi-parametric inducing point networks and neural processes.  2023\.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International conference on machine learning_ , pp. 1530–1538. PMLR, 2015.
  * Rombach et al. (2022) Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 10684–10695, 2022.
  * Roy et al. (2021) Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier.  Efficient content-based sparse attention with routing transformers.  _Transactions of the Association for Computational Linguistics_ , 9:53–68, 2021.  doi: 10.1162/tacl˙a˙00353.  URL https://aclanthology.org/2021.tacl-1.4.
  * Salimans et al. (2015) Tim Salimans, Diederik P Kingma, and Max Welling.  Markov chain monte carlo and variational inference: Bridging the gap.  In _International conference on machine learning_ , pp. 1218–1226. PMLR, 2015.
  * Si et al. (2022) Phillip Si, Allan Bishop, and Volodymyr Kuleshov.  Autoregressive quantile flows for predictive uncertainty estimation.  In _International Conference on Learning Representations_ , 2022.
  * Si et al. (2023) Phillip Si, Zeyi Chen, Subham Sekhar Sahoo, Yair Schiff, and Volodymyr Kuleshov.  Semi-autoregressive energy flows: Exploring likelihood-free training of normalizing flows.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 31732–31753. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/si23a.html.
  * Siva (2008) Nayanah Siva.  1000 genomes project.  _Nature biotechnology_ , 26(3):256–257, 2008.
  * Song et al. (2020) Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_ , 2020.
  * Tang et al. (2016) Jian Tang, Jingzhou Liu, Ming Zhang, and Qiaozhu Mei.  Visualizing large-scale and high-dimensional data.  In _Proceedings of the 25th international conference on world wide web_ , pp. 287–297, 2016.
  * Tomczak & Welling (2018) Jakub Tomczak and Max Welling.  Vae with a vampprior.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1214–1223. PMLR, 2018.
  * Vahdat & Kautz (2020) Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  _Advances in neural information processing systems_ , 33:19667–19679, 2020.
  * Vahdat et al. (2021) Arash Vahdat, Karsten Kreis, and Jan Kautz.  Score-based generative modeling in latent space.  _Advances in Neural Information Processing Systems_ , 34:11287–11302, 2021.
  * Van der Maaten & Hinton (2008) Laurens Van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(11), 2008.
  * Wang et al. (2021) Yingfan Wang, Haiyang Huang, Cynthia Rudin, and Yaron Shaposhnik.  Understanding how dimension reduction tools work: an empirical approach to deciphering t-sne, umap, trimap, and pacmap for data visualization.  _The Journal of Machine Learning Research_ , 22(1):9129–9201, 2021.
  * Wang et al. (2023a) Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov.  InfoDiffusion: Representation learning using information maximizing diffusion models.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 36336–36354. PMLR, 23–29 Jul 2023a.  URL https://proceedings.mlr.press/v202/wang23ah.html.
  * Wang et al. (2023b) Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov.  Infodiffusion: Representation learning using information maximizing diffusion models.  _arXiv preprint arXiv:2306.08757_ , 2023b.
  * Wehenkel & Louppe (2021) Antoine Wehenkel and Gilles Louppe.  Diffusion priors in variational autoencoders.  _arXiv preprint arXiv:2106.15671_ , 2021.
  * Wold et al. (1987) Svante Wold, Kim Esbensen, and Paul Geladi.  Principal component analysis.  _Chemometrics and intelligent laboratory systems_ , 2(1-3):37–52, 1987.
  * Zhang et al. (2022) Zijian Zhang, Zhou Zhao, and Zhijie Lin.  Unsupervised representation learning from pre-trained diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 35:22117–22130, 2022.
  * Zhao et al. (2017) Shengjia Zhao, Jiaming Song, and Stefano Ermon.  Infovae: Information maximizing variational autoencoders.  _arXiv preprint arXiv:1706.02262_ , 2017.

##  Appendix A Appendix

###  A.1 Computational Cost Analysis

Method NMI values at different wall-clock training times NMI @ 10 min NMI @ 20
min NMI @ 30 min NMI @ 40 min NMI @ 50 min NMI @ 60 min VAE $0.52$ $0.52$
$0.52$ $0.52$ $0.52$ $0.52$ IAF-VAE $0.54$ $0.52$ $0.52$ $0.52$ $0.52$ $0.52$
AAE $0.61$ $0.57$ $0.57$ $0.57$ $0.57$ $0.57$ DiffVAE (T=5) warm up $0.63$
$0.63$ $0.66$ $0.66$ $0.66$ DiffVAE (T=10) warm up $\mathbf{0.64}$
$\mathbf{0.68}$ $\mathbf{0.70}$ $\mathbf{0.70}$ $\mathbf{0.70}$ DiffVAE (T=20)
warm up $0.50$ $0.51$ $0.56$ $0.64$ $0.68$ DiffVAE (T=50) warm up $0.52$
$0.54$ $0.51$ $0.59$ $0.59$

Table 4: Computational cost trade-off on 1kgenome: NMI vs wall-clock training
time

We conduct a computational cost analysis between the baselines and DD-VAE with
various timesteps on the genotype clustering/visualization experiments. The
table Table 4 below shows that DD-VAE outperforms baselines at all timestamps
and continues to improve after the baselines have plateaued.

###  A.2 Priors

Below we describe the sampling process for each prior.

Pinwheel. This distribution was used in (Johnson et al., 2016). We define the
number of clusters to be 10. For semi-supervised learning experiments, this
prior is partitioned into 10 partitions, each partition being a cluster.

Swiss Roll. This distribution was used in Marsland (2014). We add noise
$\sigma=0.1$ to the prior. For semi-supervised learning experiments, this
prior is partitioned into 10 partitions. The samples from the prior, before we
add noise to it, can actually be characterized by a single scalar representing
how far you are long the swiss roll from the center. The paritioning is done
by creating 10 equal-length intervals in this 1D space.

Square. This distribution has the shaped of a square going from -1 to 1 in
both axes. Each position on the square can be characterized by a single scalar
representing how far you are from the top left corner. Sampling is done by
sampling the position uniformly and turn the 1D position to 2D latent. We add
noise $\sigma=0.06$ to the prior. For semi-supervised learning experiments,
this prior is partitioned into 10 partitions. The partitioning is done by
creating 10 equal-length intervals in the 1D position space.

VAE and IAF-VAE requires that we can evaluate the prior density. To do this,
for all priors, we evaluate the density by fitting a kernel density estimator
with mixture of gaussian kernel with bandwidth equal to 0.005, 0.008, 0.01,
0.03, and 0.05.

###  A.3 Model Architecture

All methods use the same architecture for encoder $q_{\bm{\phi}}(z|x)$ and
decoder $p_{\bm{\theta}}(x|z)$, excluding the extra parts specific to each
method which we describe below, for the same dataset. For MNIST, the encoder
and decoder are multi-layer perceptron with two hidden layers, each with 1000
hidden units. For CIFAR-10, the encoder is a 4-layer convolutional neural
network with (16, 32, 64, 128) channels with a linear layer on top, and the
decoder is a 4-layer tranposed convolutional neural network with (64, 32, 16,
3) channels where a linear layer is used to first turn the feature dimension
from 2 to 64.

IAF-VAE employs 4 IAF transformations on top of the encoder, each is
implemented with a 4-layer MADE. The number of hidden units in MADE is 24. The
ordering is reversed between every other IAF transformation.

AAE has a discriminator, used in adversarial training, which is a multi-layer
perceptron with two hidden layers, each with 1000 hidden units.

DD-VAE has a diffusion model on top of the encoder. The time-conditioned
reverse diffusion distribution is implemented with a 5-layer time-conditioned
multi-layer perceptron, each with 128 hidden units. A time-conditioned linear
layer learns an additional embedding for each timestep and adds it to the
output of the linear layer.

###  A.4 Training Details

For training, we use Adam optimizer and latent size of 2 for all of our
experiments. The training details of each algorithm are detailed below:

VAE. The batch size is set to 128. The number of epochs is 200 for
unsupervised and clustering experiments and 50 for semi-supervised
experiments. The learning rate is 0.0001. The loss is BCE for MNIST and
CIFAR-10 experiments and MSE for genotype analysis experiments. For semi-
supervised MNIST experiments, the kl divergence weight is set to be 0.01,
while for semi-supervised CIFAR-10 experiments, the kl divergence weight is
set to be 0.01. For other experiments, the KL divergence weight is set with a
schedule linear on number of epochs going from 0 to 0.01. We also have a
weight of 5 multiplied to the prior density.

IAF-VAE. The batch size, number of epochs, learning rate, loss, KL divergence
weight, and prior density weight are the same as VAE. The context size, i.e.,
the size of features used to initialize the flow layers for different datat
point, is 10.

AAE. The batch size is set to 128. The number of epochs is 200 for all
experiments. The learning rate is 0.0002. The loss is MSE for all experiments.
To stabilize the training, we add noise to the input to the discriminator with
sigma 0.3 at the start and lower it by 0.1 for every 50 epochs. The noise
equals to 0 at epoch 150.

DD-VAE. The batch size is set to 128 for most experiments, except for semi-
supervied experiments where the batch size is 1024. The number of epochs is
200 for unsupervised and clustering experiments and 30 for semi-supervised
experiments. The learning rate is 0.0001. The loss is BCE for MNIST and
CIFAR-10 experiments and MSE for genotype analysis experiments. For
unsupervised MNIST and CIFAR-10 experiments, the KL divergence weight is set
to 0.003. For semi-supervised MNIST experiment, we use KL divergence weight of
0.1. For semi-supervised CIFAR-10 experiment, we use KL divergence weight of
0.5. For clustering experiment, we use KL divergence weight of 0.005. The
number of timesteps is 20 for unsupervised and clustering experiments and 100
for semi-supervised experiments.

###  A.5 Genotype Analysis Experiments Details

Figure 5: Legend showing what ethnicity each color corresponds to in the 1000
Genomes dataset

Before inputting the data points into any of the visualization methods, we
first pre-process it by running a PCA and keep only the first 1000 principal
components of the data points. We further divide the features by 30 for all
latent variables model methods.

The legend of the 1000 Genomes Visualization plot can be found at Figure 5.

Method Pinwheel Swiss Roll Square Acc Latent NLL Acc Latent NLL Acc Latent NLL
VAE $0.16\pm 0.01$ $86.9\pm 99.03$ $0.15\pm 0.01$ $360.41\pm 185.87$ $0.15\pm
0.01$ $63.3\pm 70.64$ IAF-VAE $0.19\pm 0.01$ $3.12\pm 1.16$ $0.16\pm 0.01$
$189.39\pm 57.48$ $0.18\pm 0.01$ $0.99\pm 0.11$ AAE $0.23\pm 0.00$ $2.30\pm
0.07$ $\bm{0.23\pm 0.01}$ $\bm{2.84\pm 0.08}$ $0.20\pm 0.01$ $1.46\pm 0.49$
HVAE w/ flow $0.20\pm 0.01$ $2.06\pm 0.13$ $0.20\pm 0.01$ $106.15\pm 15.30$
$0.19\pm 0.01$ $1.05\pm 0.06$ DiffVAE $\bm{0.23\pm 0.01}$ $\bm{1.37\pm 0.02}$
$\bm{0.23\pm 0.01}$ $\bm{2.80\pm 0.11}$ $\bm{0.24\pm 0.01}$ $\bm{0.86\pm
0.05}$

Table 5: Unsupervised learning on CIFAR-10. We report accuracy using KNN
(K=20) classifier and latent negative log-likelihood with pinwheel, swiss
roll, and square priors.

Method Pinwheel Swiss Roll Square Latent NLL Avg Class Latent NLL Avg Class
Latent NLL Avg Class Latent NLL Latent NLL Latent NLL VAE $7.67\pm 1.34$
$9.84\pm 2.87$ $49.64\pm 28.12$ $128.46\pm 47.08$ $1.51\pm 0.23$ $0.65\pm
0.37$ IAF-VAE $4.54\pm 0.59$ $4.46\pm 0.99$ $64.27\pm 73.87$ $102.16\pm
106.36$ $1.04\pm 0.07$ $0.07\pm 0.09$ AAE $2.33\pm 0.02$ $2.42\pm 0.06$
$3.76\pm 0.04$ $4.23\pm 0.13$ $1.45\pm 0.12$ $1.31\pm 0.20$ DiffVAE
$\bm{1.49\pm 0.09}$ $\bm{0.35\pm 0.18}$ $\bm{3.35\pm 0.27}$ $\bm{2.28\pm
0.39}$ $\bm{0.97\pm 0.03}$ $\bm{-0.29\pm 0.04}$

Table 6: Semi-supervised learning on CIFAR-10 (10,000 labels). We report
average class latent negative log-likelihood and latent negative log-
likelihood with pinwheel, swiss roll, and square priors.

Method Latent NLL - Pinwheel Latent NLL - Swiss Roll Latent NLL - Square VAE
$6.18\pm 0.93$ $82.96\pm 99.30$ $25.07\pm 62.41$ IAF-VAE $2.06\pm 0.21$
$12.13\pm 18.28$ $1.53\pm 0.82$ AAE $1.94\pm 0.09$ $3.43\pm 0.14$ $1.97\pm
0.70$ HVAE w/ flow $2.18\pm 0.13$ $4.17\pm 1.23$ $\mathbf{1.25\pm 0.08}$
DiffVAE $\mathbf{1.57\pm 0.10}$ $\mathbf{3.19\pm 0.07}$ $\mathbf{1.24\pm
0.13}$ VAE with Diffusion (standard ELBO) $6.09$ $510.53$ $265.25$

Table 7: Unsupervised learning on MNIST, including the failed results of VAE
with Diffusion

Model Training Objective Approximating Family Sample-based Prior Auxiliary
Variable Tasks Simplified Graphical Illustration VAE ELBO Diagonal Gaussian No
No Density estimation $\mathbf{x}\rightarrow\mathbf{z}\rightarrow\mathbf{x}$
IAF-VAE ELBO Normalizing flow No Yes Density estimation / Visualization
$\mathbf{x}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{z}_{T}\rightarrow\mathbf{x}$
AAE Adversarial training Adversarial generator Yes No Visualization
$\mathbf{x}\rightarrow\mathbf{z}\rightarrow\mathbf{x}$ HVAE w/o and w/ flow
ELBO Factorial Normal / Normalizing flow No Yes Density estimation / High-
quality sample generation
$\mathbf{x}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{z}_{T}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{x}$
ADGM ELBO Non-Gaussian No Yes Density estimation
$\mathbf{x}\rightarrow\mathbf{a}\rightarrow\mathbf{z}\rightarrow\mathbf{x}$
LDM ELBO Diagonal Gaussian No Yes High-quality sample generation
$\mathbf{x}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{z}_{T}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{x}$
LSGM ELBO+score matching Diagonal Gaussian No Yes High-quality sample
generation
$\mathbf{x}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{z}_{T}\rightarrow\mathbf{z}_{0}\rightarrow\mathbf{x}$
DiffVAE ELBO+sleep term Denoising diffusion Yes Yes Density estimation /
Visualization
$\mathbf{x}\rightarrow\mathbf{z}_{T}(\mathbf{y})\rightarrow\mathbf{z}_{0}(\mathbf{z})\rightarrow\mathbf{x}$

Table 8: Comparison of DiffVAE to other relevant methods