Variational Autoencoder

Introduction

In my previous article “Expectation Maximization Algorithm”, we discussed how to optimize probabilistic models with latent variables using the Expectation Maximization (EM) algorithm.

In this article, I would like to discuss how to optimize probabilistic models with latent variables using Variational Autoencoder (VAE).

Variational Autoencoder

Similar to the maximization - maximization procedure interpretation for the EM algorithm, the log evidence of a probabilistic model consisting of parameters $\mathit{\theta }$, $\log p_{\mathit{\theta }}(\mathbf{x})$, can be expressed as the sum of an evidence lower bound (ELBO) and a Kullback-Leibler (KL) divergence.

$$\begin{array}{rl}\log p_{\mathit{\theta }}(\mathbf{x})& ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log p_{\mathit{\theta }}(\mathbf{x})]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})q(\mathbf{z})}{q(\mathbf{z})p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q(\mathbf{z})}+\log \frac{q(\mathbf{z})}{p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q(\mathbf{z})}]+{\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{q(\mathbf{z})}{p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q(\mathbf{z})}]-{\mathbb{E}}_{\mathbf{z}\sim q(\mathbf{z})}[\log \frac{p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}{q(\mathbf{z})}]\\ & =\text{ELBO}(q,\mathit{\theta })+D_{\text{KL}}(q(\mathbf{z})||p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x}))\end{array}$$

where $\text{ELBO}(q,\mathit{\theta })$ is the evidence lower bound (ELBO) and $D_{\text{KL}}(q(\mathbf{z})||p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x}))$ is the Kullback-Leibler (KL) divergence between the approximate posterior $q(\mathbf{z})$ and the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$.

In the iterative EM algorithm, the approximate posterior $q(\mathbf{z})$ is not learned explicitly. In the VAE, the approximate posterior $q(\mathbf{z})$ is learned explicitly using an inference model. Specifically, the approximate posterior $q(\mathbf{z})$ is parameterized by a neural network, $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$, where $\mathit{\varphi }$ are the variational parameters of the neural network ($\mathit{\varphi }$ are called variational parameters because they are the parameters of the variational distribution $q(\mathbf{z}|\mathbf{x})$), conditioned on the observed data $\mathbf{x}$. Concretely,

$$\begin{array}{rl}\log p_{\mathit{\theta }}(\mathbf{x})& ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}]+{\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log \frac{q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}{p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})}]\\ & =\text{ELBO}(\mathit{\varphi },\mathit{\theta })+D_{\text{KL}}(q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})||p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x}))\end{array}$$

where $\text{ELBO}(\mathit{\varphi },\mathit{\theta })$ is the evidence lower bound (ELBO) and $D_{\text{KL}}(q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})||p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x}))$ is the Kullback-Leibler (KL) divergence between the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ and the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$.

Model Parameters Optimization

In the $\text{ELBO}(\mathit{\varphi },\mathit{\theta })$, $p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})$ is called the generative model and $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ is called the inference model. $p_{\mathit{\theta }}(\mathbf{x})$ is also the generative model, even though it is not explicitly modeled.

$D_{\text{KL}}(q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})||p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x}))$ is non-negative and it is zero if and only if the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ is equal to the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$.

If the parameters $\mathit{\theta }$ are fixed and only the parameters $\mathit{\varphi }$ are optimized, maximizing the ELBO is equivalent to minimizing the KL divergence between the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ and the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$, making the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ as close as possible to the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$.

If the parameters $\mathit{\varphi }$ are fixed and only the parameters $\mathit{\theta }$ are optimized, maximizing the ELBO is approximately equivalent to maximizing the log evidence $\log p_{\mathit{\theta }}(\mathbf{x})$, making the generative model $p_{\mathit{\theta }}(\mathbf{x})$ better.

Therefore, the parameters $\mathit{\varphi }$ and $\mathit{\theta }$ can even be jointly optimized by maximizing the ELBO so that the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ becomes close to the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$ and the generative model $p_{\mathit{\theta }}(\mathbf{x})$ becomes better.

Gradient-based optimization algorithms can be used to optimize the parameters $\mathit{\varphi }$ and $\mathit{\theta }$. The target function is defined as follows.

$$\begin{array}{rl}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& =\text{ELBO}(\mathit{\varphi },\mathit{\theta })\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log \frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\end{array}$$

The gradients of the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ with respect to the generative model parameters $\mathit{\theta }$ are straightforward to obtain.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\mathit{\theta }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathrm{\nabla }}_{\mathit{\theta }}{\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}\left[{\mathrm{\nabla }}_{\mathit{\theta }}(\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x}))\right]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[{\mathrm{\nabla }}_{\mathit{\theta }}\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-{\mathrm{\nabla }}_{\mathit{\theta }}\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[{\mathrm{\nabla }}_{\mathit{\theta }}\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})]\end{array}$$

This ${\mathrm{\nabla }}_{\mathit{\theta }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ can be approximated using an unbiased Monte Carlo estimate by sampling ${\mathbf{z}}^{(i)}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\mathit{\theta }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& \approx \frac{1}{M}\sum _{i=1}^M{\mathrm{\nabla }}_{\mathit{\theta }}\log p_{\mathit{\theta }}(\mathbf{x},{\mathbf{z}}^{(i)})\end{array}$$

where $M$ is the number of samples which can be just 1.

The gradients of the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ with respect to the inference model parameters $\mathit{\varphi }$ are somewhat more complicated to obtain.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\mathit{\varphi }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathrm{\nabla }}_{\mathit{\varphi }}{\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\end{array}$$

Here we encountered the same problem as we did in the EM algorithm. The probability distribution that the latent variables $\mathbf{z}$ follow is parameterized by $\mathit{\varphi }$, and the parameters to be optimized are also $\mathit{\varphi }$. Therefore, the Monte Carlo estimate of ${\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]$ by sampling ${\mathbf{z}}^{(i)}\sim q_{{\mathit{\varphi }}^{\text{old}}}(\mathbf{z}|\mathbf{x})$ is biased. In addition, even if the Monte Carlo estimate is unbiased, the sampling process is not differentiable with respect to $\mathit{\varphi }$ and the gradient methods would become invalid.

If the latent variables are transformed from some other latent variables that are not parameterized by $\mathit{\varphi }$ via a deterministic transformation, the unbiased Monte Carlo estimate can be obtained and the sampling process becomes differentiable with respect to $\mathit{\varphi }$. This is the idea of the reparameterization trick.

Concretely, we have some differentiable and invertible transformation $g$ that transforms some other latent variables $\mathit{ϵ}$ that are not parameterized by $\mathit{\varphi }$ into the latent variables $\mathbf{z}$ that are parameterized by $\mathit{\varphi }$.

$$\begin{array}{rl}\mathbf{z}& =g(\mathit{ϵ},\mathit{\varphi },\mathbf{x})\\ \mathit{ϵ}& \sim p(\mathit{ϵ})\end{array}$$

The distribution of $\mathit{ϵ}$ is usually chosen to be a distribution $p(\mathit{ϵ})$ such that the distribution of $\mathbf{z}$ is the same as $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$.

As a result, the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ can be expressed as an expectation over the latent variables $\mathit{ϵ}$ that are not parameterized by $\mathit{\varphi }$ and this target function is completely differentiable with respect to $\mathit{\theta }$ and $\mathit{\varphi }$.

$$\begin{array}{rl}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\end{array}$$

where $\mathbf{z}=g(\mathit{ϵ},\mathit{\varphi },\mathbf{x})$.

The gradients of the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ with respect to the generative model parameters $\mathit{\theta }$ and the inference model parameters $\mathit{\varphi }$ become straightforward to obtain.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\mathit{\theta }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathrm{\nabla }}_{\mathit{\theta }}{\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}\left[{\mathrm{\nabla }}_{\mathit{\theta }}(\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x}))\right]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[{\mathrm{\nabla }}_{\mathit{\theta }}\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})]\end{array}$$

$$\begin{array}{rl}{\mathrm{\nabla }}_{\mathit{\varphi }}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathrm{\nabla }}_{\mathit{\varphi }}{\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}\left[{\mathrm{\nabla }}_{\mathit{\varphi }}(\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x}))\right]\end{array}$$

According to the reparameterization trick and the transformations of random variables, the probability distribution $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ can be transformed from the probability distribution $p(\mathit{ϵ})$.

$$\begin{array}{rl}{q}_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})& =\frac{p(\mathit{ϵ})}{|det{\mathbf{J}}_g(\mathit{ϵ})|}\end{array}$$

where ${\mathbf{J}}_g(\mathit{ϵ})$ is the Jacobian matrix of the transformation $g$ at $\mathit{ϵ}$.

$$\begin{array}{rl}{\mathbf{J}}_g(\mathit{ϵ})& =\frac{\partial \mathbf{z}}{\partial \mathit{ϵ}}\\ & =\left[\begin{array}{cccc}\frac{\partial z_1}{\partial {ϵ}_1}& \frac{\partial z_1}{\partial {ϵ}_2}& \cdots & \frac{\partial z_1}{\partial {ϵ}_n}\\ \frac{\partial z_2}{\partial {ϵ}_1}& \frac{\partial z_2}{\partial {ϵ}_2}& \cdots & \frac{\partial z_2}{\partial {ϵ}_n}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{\partial z_n}{\partial {ϵ}_1}& \frac{\partial z_n}{\partial {ϵ}_2}& \cdots & \frac{\partial z_n}{\partial {ϵ}_n}\end{array}\right]\end{array}$$

Thus, the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ can be rewritten as follows.

$$\begin{array}{rl}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log \frac{p(\mathit{ϵ})}{|det{\mathbf{J}}_g(\mathit{ϵ})|}]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})-\log p(\mathit{ϵ})+\log |det{\mathbf{J}}_g(\mathit{ϵ})|]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log (p_{\mathit{\theta }}(\mathbf{x}|\mathbf{z})p_{\mathit{\theta }}(\mathbf{z}))-\log p(\mathit{ϵ})+\log |det{\mathbf{J}}_g(\mathit{ϵ})|]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x}|\mathbf{z})+\log p_{\mathit{\theta }}(\mathbf{z})-\log p(\mathit{ϵ})+\log |det{\mathbf{J}}_g(\mathit{ϵ})|]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x}|g(\mathit{ϵ},\mathit{\varphi },\mathbf{x}))+\log p_{\mathit{\theta }}(g(\mathit{ϵ},\mathit{\varphi },\mathbf{x}))-\log p(\mathit{ϵ})+\log |det{\mathbf{J}}_g(\mathit{ϵ})|]\\ & ={\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[\log p_{\mathit{\theta }}(\mathbf{x}|g(\mathit{ϵ},\mathit{\varphi },\mathbf{x}))+\log p_{\mathit{\theta }}(g(\mathit{ϵ},\mathit{\varphi },\mathbf{x}))+\log |det{\mathbf{J}}_g(\mathit{ϵ})|]+\text{constant}\end{array}$$

where $\text{constant}$ is just ${\mathbb{E}}_{\mathit{ϵ}\sim p(\mathit{ϵ})}[-\log p(\mathit{ϵ})]$ which does not depend on $\mathit{\varphi }$ and $\mathit{\theta }$.

The gradients of the target function ${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})$ with respect to the generative model parameters $\mathit{\theta }$ and the inference model parameters $\mathit{\varphi }$ can be approximated using an unbiased Monte Carlo estimate by sampling ${\mathit{ϵ}}^{(i)}\sim p(\mathit{ϵ})$ and easily computed in automatic differentiation frameworks, such as PyTorch and TensorFlow.

Generative Model Evaluation

To evaluate the performance of the generative model $p_{\mathit{\theta }}(\mathbf{x})$, the marginal likelihood $\log p_{\mathit{\theta }}(\mathbf{x})$ can be estimated via importance sampling.

$$\begin{array}{rl}\log p_{\mathit{\theta }}(\mathbf{x})& =\log {\mathbb{E}}_{\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}\left[\frac{p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})}{q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})}\right]\\ & \approx \log \frac{1}{M}\sum _{i=1}^M\frac{p_{\mathit{\theta }}(\mathbf{x},{\mathbf{z}}^{(i)})}{q_{\mathit{\varphi }}({\mathbf{z}}^{(i)}|\mathbf{x})}\end{array}$$

where ${\mathbf{z}}^{(i)}\sim q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ and $M$ is the number of samples.

One can also sample from the generative model $p_{\mathit{\theta }}(\mathbf{x},\mathbf{z})$ and evaluate the quality of the generated sample $\mathbf{x}$ using some other methods, such as visual inspection. We will first sample the latent variables $\mathbf{z}$ from the prior distribution $p_{\mathit{\theta }}(\mathbf{z})$ and then sample the observed data $\mathbf{x}$ from the generative model $p_{\mathit{\theta }}(\mathbf{x}|\mathbf{z})$.

Increasing the Inference Model Flexibility

What if the inference model $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ is not flexible enough to approximate the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$?

The tightness of the ELBO depends on the flexibility of the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$. If $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ is not flexible enough to approximate the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$, the ELBO may not be tight enough to approximate the log evidence $\log p_{\mathit{\theta }}(\mathbf{x})$ well. In this case, the generative model $p_{\mathit{\theta }}(\mathbf{x})$ may not be learned well.

In many simple settings, the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ is just a simple multivariate Gaussian distribution, and it is insufficient to approximate the true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$ if true posterior $p_{\mathit{\theta }}(\mathbf{z}|\mathbf{x})$ is more complex.

Similar to what we have discussed in the article “Expectation Maximization Algorithm”, having more hierarchical latent variables can make the approximate posterior $q_{\mathit{\varphi }}(\mathbf{z}|\mathbf{x})$ more flexible and the ELBO tighter. In our case, we will introduce more hierarchical latent variables to both the generative model and the inference model.

Suppose we have $n$ sets of latent variables ${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n$ and the generative model $p_{\mathit{\theta }}(\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)$ and the inference model $q_{\mathit{\varphi }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n|\mathbf{x})$. The ELBO can be expressed as follows.

$$\begin{array}{rl}\text{ELBO}(\mathit{\varphi },\mathit{\theta })& ={\mathbb{E}}_{{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n\sim q_{\mathit{\varphi }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n|\mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)-\log q_{\mathit{\varphi }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n|\mathbf{x})]\end{array}$$

In the generative model, we can introduce more hierarchical latent variables by factorizing the joint distribution of the observed data $\mathbf{x}$ and the latent variables ${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n$.

$$\begin{array}{rl}{p}_{\mathit{\theta }}(\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)& =p_{\mathit{\theta }}(\mathbf{x}|{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)p_{\mathit{\theta }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)\\ & =p_{\mathit{\theta }}(\mathbf{x}|{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)p_{\mathit{\theta }}({\mathbf{z}}_1|{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)p_{\mathit{\theta }}({\mathbf{z}}_2,{\mathbf{z}}_3,\cdots ,{\mathbf{z}}_n)\\ & ⋮\\ & =p_{\mathit{\theta }}(\mathbf{x}|{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)p_{\mathit{\theta }}({\mathbf{z}}_1|{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)p_{\mathit{\theta }}({\mathbf{z}}_2|{\mathbf{z}}_3,\cdots ,{\mathbf{z}}_n)\cdots p_{\mathit{\theta }}({\mathbf{z}}_n)\end{array}$$

Given the generative model, there can be different ways to introduce more hierarchical latent variables to the inference model.

One common way is to follow the same hierarchy as the generative model.

$$\begin{array}{rl}{q}_{\mathit{\varphi }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n|\mathbf{x})& =q_{\mathit{\varphi }}({\mathbf{z}}_1|{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n,\mathbf{x})q_{\mathit{\varphi }}({\mathbf{z}}_2|{\mathbf{z}}_3,\cdots ,{\mathbf{z}}_n,\mathbf{x})\cdots q_{\mathit{\varphi }}({\mathbf{z}}_n|\mathbf{x})\end{array}$$

In the first step, the inference model will sample the latent variables ${\mathbf{z}}_n$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_n|\mathbf{x})$, and the sampled latent variables ${\mathbf{z}}_n$ can be used to compute the prior $p_{\mathit{\theta }}({\mathbf{z}}_n)$.

In the second step, the inference model will sample the latent variables ${\mathbf{z}}_{n-1}$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_{n-1}|{\mathbf{z}}_n,\mathbf{x})$, and the sampled latent variables ${\mathbf{z}}_{n-1}$ and ${\mathbf{z}}_n$ can be used to compute the prior $p_{\mathit{\theta }}({\mathbf{z}}_{n-1}|{\mathbf{z}}_n)$.

In the $n$-th step, the inference model will sample the latent variables ${\mathbf{z}}_1$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_1|{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n,\mathbf{x})$, and the sampled latent variables ${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n$ can be used to compute the prior $p_{\mathit{\theta }}({\mathbf{z}}_1|{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)$.

This kind of inference model is sometimes also called the top-down inference model.

The other common way is reverse the hierarchy of the generative model.

$$\begin{array}{rl}{q}_{\mathit{\varphi }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n|\mathbf{x})& =q_{\mathit{\varphi }}({\mathbf{z}}_n|\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_{n-1})q_{\mathit{\varphi }}({\mathbf{z}}_{n-1}|\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_{n-2})\cdots q_{\mathit{\varphi }}({\mathbf{z}}_1|\mathbf{x})\end{array}$$

In the first step, the inference model will sample the latent variables ${\mathbf{z}}_1$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_1|\mathbf{x})$.

In the second step, the inference model will sample the latent variables ${\mathbf{z}}_2$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_2|\mathbf{x},{\mathbf{z}}_1)$.

In the $n$-th step, the inference model will sample the latent variables ${\mathbf{z}}_n$ from the approximate posterior $q_{\mathit{\varphi }}({\mathbf{z}}_n|\mathbf{x},{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_{n-1})$.

Once all the latent variables ${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n$ are sampled, the sampled latent variables can be used to compute the prior $p_{\mathit{\theta }}({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_n)$ in the generative model.

This kind of inference model is sometimes also called the bottom-up inference model.

It seems that the top-down inference model is more favorable than the bottom-up inference model in practice.

References

• An Introduction to Variational Autoencoders
• Improving Variational Inference with Inverse Autoregressive Flow
• Expectation Maximization Algorithm
• Reparameterization Trick
• Transformations of Random Variables