VAEs: Using Neural Networks To Approximate Conditional Distributions

Question

Given the setup of a VAE, such as that outlined in Kingma and Welling (2014), there is a conditional probability distribution $p(z|x)$ describing the distribution of generated data $z \in \mathcal{Z}$ given sample data $x \in \mathcal{X}$, and there is likewise a conditional probability distribution $p(x|z)$ for the reverse. We do not know either of these distributions, so we approximate them using neural networks $q_{\phi}(z|x)$ and $p_{\theta}(x|z)$ with parameters $\phi$ and $\theta$. We then find $\phi$ and $\theta$ such that the ELBO (Evidence Lower BOund) is minimized.

Here is what confuses me:

$p(z|x)$ and $p(x|z)$ are conditional distributions. This means that for a fixed $x \in \mathcal{X}$, $\displaystyle \int_{\mathcal{Z}} p(z|x) \; dz = 1$, and similarly, for a fixed $z \in \mathcal{Z}$, $\displaystyle \int_{\mathcal{X}} p(x|z) \; dx = 1$. But we are approximating them using neural networks $q_{\phi}(z|x)$ and $p_{\theta}(x|z)$, which are in general just continuous functions parametrized by $\phi$ and $\theta$. So why is this a valid approximation when there is no requirement for a neural network to integrate to $1$ like $p(z|x)$ and $p(x|z)$ do?

Answer 1

When we say that $q_\phi(z|x)$ "is" a neural network, what we actually mean is that $q_\phi$ is a multivariate normal distribution, and we have a neural network $g$ which outputs $\mu, \sigma = g(x)$ (alternatively you might say there are two neural networks $\mu(x)$ and $\sigma(x)$, so that $q_\phi(z|x) = \mathcal{N}(\mu, \sigma^2 I)$. And of course this is always a valid distribution. The same holds for $p_\theta(x|z)$.

Answer 2

In case of VAE, both inference network p(z/x) and generative network p(x/z) output a gaussian distribution parameter which are sampled. So the latent space is sampled from p(z/x) and the output (generated image as an example) is sampled from p(x/z) gaussian distribution