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

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

• But we can compute the KL divergence between N(0,I) and Q(z|X) in closed form if both are Gaussian, which is what we typically do. Once trained, the distribution of Q(z|X) will look like a quilt of 10 Gaussians for MNIST (analogous to the picture you described for MNIST on the uniform distribution), and the mixture will be bigger than N(0,I). I guess we have to manually inspect this mixture distribution in the latent space and try to come up with a way to sample it. On the other hand, this blog claims that we should sample from N(0,I) at generator time, not Q(z|X). – foghorn Mar 12 '18 at 23:17
• the blog: towardsdatascience.com/… – foghorn Mar 12 '18 at 23:18
• I edited my answer to answer your question on why you might expect the "quilt" to be approximately N(0,I) – shimao Mar 13 '18 at 18:25
• @shimao I think it is misleading to say $KL(Q(z|x) \Vert P(z))$ is a regularizer term in some tutorials, it is simply a term in the lower bound. Moreover, for generative purposes, the only thing that we care about is $P(x) = \int P(x|z)P(z)$. $P(z)$ is the prior and we get $P(x|z)$ by training, so we don't actually care about $Q(z|x)$. – meTchaikovsky May 12 '18 at 5:23
• @me_Tchaikovsky i'm not exactly sure what you're getting at -- but there is indeed value in knowing how close $Q(z|x)$ comes to $P(z|x)$, because that determines whether it makes sense to sample from VAEs in the straightforward way (sampling the latent space and then running the decoder) – shimao May 12 '18 at 5:28

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