How does the VAE learn a joint distribution?

Question

I found the following paragraph from An Introduction to Variational Autoencoders sounds relevant, but I am not fully understanding it.

A VAE learns stochastic mappings between an observed $\mathbf{x}$-space, whose empirical distribution $q_{\mathcal{D}}(\mathbf{x})$ is typically complicated, and a latent $\mathbf{z}$-space, whose distribution can be relatively simple (such as spherical, as in this figure). The generative model learns a joint distribution $p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})$ that is often (but not always) factorized as $p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})=p_{\boldsymbol{\theta}}(\mathbf{z}) p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$, with a prior distribution over latent space $p_{\boldsymbol{\theta}}(\mathbf{z})$, and a stochastic decoder $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$. The stochastic encoder $q_{\phi}(\mathbf{z} \mid \mathbf{x})$, also called inference model, approximates the true but intractable posterior $p_{\theta}(\mathbf{z} \mid \mathbf{x})$ of the generative model.

How is it that the generative model learns a joint distribution $p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})$ in the case of the VAE? I know that learning the weights of the decoder is learning $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$

Answer 1

The VAE models the following directed graphical model (figure 1 from the original VAE paper)

So, you have 2 sets of parameters, $\boldsymbol{\phi}$ and $\boldsymbol{\theta}$, and 2 random variables, $\mathbf{z}$ (latent r.v.) and $\mathbf{x}$.

How can you view this graphical model (in figure 1 above) as a generative model?

1. First, a sample $\mathbf{z}^{(i)}$ is generated from the (variational) probability distribution $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$

2. Then, a sample $\mathbf{x}^{(i)}$ is generated from $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}^{(i)})$

Now, more concretely, let's assume that we have a dataset $\mathbf{X}=\left\{\mathbf{x}^{(i)}\right\}_{i=1}^{N}$ of $N$ i.i.d. samples of the random variable $\mathbf{x}$. So, each of these $\mathbf{x}^{(i)}$ has been generated as follows

1. a sample $\mathbf{z}^{(i)}$ is generated from some prior distribution $p_{\boldsymbol{\theta}}(\mathbf{z})$

2. a sample $\mathbf{x}^{(i)}$ is generated from some conditional distribution $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}^{(i)})$.

Now, in the VAE, the encoder represents $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$, while the decoder represents $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$. If you want to train a VAE, you also need to make an assumption about $p_{\boldsymbol{\theta}}(\mathbf{z})$, for example, you can assume $p_{\boldsymbol{\theta}}(\mathbf{z}) = \mathcal{N}(\mathbf{z} ; \mathbf{0}, \mathbf{I})$. Once trained, you use the variational distribution (encoder) as the prior from which you sample $\mathbf{z}^{(i)}$ (although we train the VAE as if the variational distribution is an approximation of the true/unknown/intractable posterior given a usually fixed prior and the likelihood/decoder), in order to sample $\mathbf{x}^{(i)}$. This is not wrong. In fact, this is just how you usually do Bayesian statistics. You have a prior and a likelihood (and maybe a marginal), then you compute the posterior, which then becomes the new prior. So, if you had more data, you could learn a new variational distribution $q'_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$ by assuming that your new prior is $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$.

If you keep in mind the following equation

$$p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})=p_{\boldsymbol{\theta}}(\mathbf{z}) p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$$

it should remind you why the VAE can be viewed as a generative model.