2
$\begingroup$

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. enter image description here

|cite|improve this question
$\endgroup$
2
$\begingroup$

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

|cite|improve this answer
$\endgroup$
  • $\begingroup$ 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). $\endgroup$ – foghorn Mar 12 '18 at 23:17
  • $\begingroup$ the blog: towardsdatascience.com/… $\endgroup$ – foghorn Mar 12 '18 at 23:18
  • $\begingroup$ I edited my answer to answer your question on why you might expect the "quilt" to be approximately N(0,I) $\endgroup$ – shimao Mar 13 '18 at 18:25
  • $\begingroup$ @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)$. $\endgroup$ – meTchaikovsky May 12 '18 at 5:23
  • $\begingroup$ @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) $\endgroup$ – shimao May 12 '18 at 5:28
1
$\begingroup$

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

|cite|improve this answer
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.