How well does $Q(z|X)$ match $N(0,I)$ in variational autoencoders?

Question

At train time, the KL divergence term drives $Q(z=\mu(X)+\epsilon \times\Sigma(X) | X)$ toward $N(0,I)$, where $\epsilon\sim N(0,I)$. It can't drive $Q(z|X)$ to exactly $N(0,I)$ because the reconstruction loss of the encoder/decoder pair would explode (the $Q(z|X)$ network would destroy all information about $X$).

Therefore when we run the system at "generator time" using only the decoder and sampling $z$ from $N(0,I)$, won't this poorly represent the training set because $Q(z|X)$ over the training set is too different from $N(0,I)$? For example $Q(z|X)$ might look like $N(0,2\times I)$, or it might even have some nonlinear hard-to-sample shape.

edit1: To clarify and ask a more well defined question: If the distribution of Q(z|X) is significantly different from N(0,I), why do we sample from N(0,I) when generating samples? Won't this yield samples that poorly represent the training set?

edit2: Even more clarification. This image shows the 10 MNIST digits mapped into a 2D latent space. You can see it does not match $\mathcal{N}(0,I)$. This image is based on 2 latent dimensions and 2 hidden layer encoder, each with 500 nodes.

Answer 1

It's ok for $Q(z|X)$ to be different from $\mathcal{N}(0, I)$, because when we sample from the VAE, we're not trying to reconstruct $X$ anymore. Instead, we're trying to sample some $X \sim \mathcal{X}$ where $\mathcal{X}$ is the distribution of all images in the dataset.

Imagine of the latent space were actually a uniform distribution over the interval $(0,10)$, and we were autoencoding MNIST digits. Suppose that images with 1 in them happened to have $Q(z|X)$ distributed around $(0,1)$, images with 2 happened to be around $(1,2)$, etc.

Then for any particular $X$, $Q(z|X)$ is not close to matching the uniform distribution. However, as long as the mixture $\frac{1}{n} \sum_i Q(z|X_i)$ reasonably covers and matches the uniform distribution, it's reasonable to sample $z \sim U(0,10)$ and then run the decoder, because the $z$ you got is probably close to $\mu(X)$ for some $X$.

edit: To answer the question of why we might expect the mixture of $Q(z|X)$ to be approximately $\mathcal{N}(0,I)$, note that we can decompose $P(z) = \int P(z|X) p(X) dz = E\left[ P(z|X) \right]$. By definition, $z \sim \mathcal{N}(0,I)$. However, when we approximate $P(z|X)$ with the encoder $Q(z|X)$, we end up with something slightly different.

Minimizing the VAE loss is equivalent to maximizing $\log P(X) - \mathcal{D}_\text{KL}(Q(z|X) || P(z|X))$. So we're simultaneously maximizing the log likelihood of the data while also encouraging $Q(z|X)$ to be as close to $P(z|X)$ as possible. As a result, we should end up with very close to $\mathcal{N}(0,I)$.

Answer 2

The answer is two-fold:

(1) The encoder network needs to be expressive enough (wide enough and deep enough) to be able to map the nonlinear input space to something close to $\mathcal{N}(0,I)$.

(2) In addition to (1) (I added a 3rd hidden layer to the MNIST example I described in the question), when I increase the number of latent dimensions, I observe that the mapping of the training data into the latent space becomes closer to $\mathcal{N}(0,I)$. In hindsight this is not super surprising, because the system is able to store information across more dimensions, so each individual latent dimension can get closer to $\mathcal{N}(0,I)$.

Answer 3

Also, we sample $z$ from the prior and decode it because we made the assumption that $p_{\theta}(x) = \int p_{\theta}(x|z)p(z) dz$. Therefore, sampling from $p_{\theta}(x)$ is equivalent to sampling from the joint $p_{\theta}(x,z)$ then discarding $z$.