# How does the VAE learn a joint distribution?

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})$$

## 1 Answer

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.

• Thanks, this is a great explanation. I was wondering if we considered the variational posterior as a prior, but wasn't sure since it wasn't stated anywhere explicitly, so now that is cleared up for me. Oct 29, 2021 at 18:38