Why do we penalize individual example divergence in variational autoencoder?

Question

In variational autoencoder, we want to learn a mapping between input space $X$ and latent space $Z$, and $z\in Z$ is related to $x\in X$ with $z\sim MVN(\mu(x), \Sigma(x))$. In addition, we desire that the latent vectors $z$ end up being distributed as a unit multivariate Normal.

To train a VAE, we penalize, for each example $x$, the deviation of $MVN(\mu(x), \Sigma(x))$ to the unit MVN. Intuitively, this doesn't feel so right, because we want the latents of general population to be jointly distributed as a unit MVN.

As a specific example, if I am encoding the MNIST digit dataset, if I encode all possible digits, I would expect it to be distributed as a unit MVN. But if I select only digit 1s, and then encode them, we shouldn't want those latents to be a unit MVN right? For example, if it occupies ``a half'' of the unit MVN ball, maybe latents of digit 8s occupy the other half, and in the end, over all digits, the latents still look like a unit MVN. Thus, it feels a bit weird to me that the objective is regularizing divergence of individual mean and variance to that of a unit MVN. What am I missing here?

Answer 1

$z\in Z$ is related to $x\in X$ with $z\sim MVN(\mu(x), \Sigma(x))$.

Almost but not quite -- $q(z|x)$ is $\mathcal{N}(\mu(x), \Sigma(x))$.

In addition, we desire that the latent vectors $z$ end up being distributed as a unit multivariate Normal.

To be pedantic, the model defines $z$ to be distributed as a standard multivariate normal, which is stronger than merely "desiring".

As a specific example, if I am encoding the MNIST digit dataset, if I encode all possible digits, I would expect it to be distributed as a unit MVN. But if I select only digit 1s, and then encode them, we shouldn't want those latents to be a unit MVN right?

If you select only 1s then you will learn a distribution over (images of) 1s, so why shouldn't they spread themselves over the latent "ball"?

Thus, it feels a bit weird to me that the objective is regularizing divergence of individual mean and variance to that of a unit MVN. What am I missing here?

I think it is misleading at best to think interpret the VAE objective as reconstruction + regularization/penalty term. It's crucial to understand the point of variational inference and the derivation of the evidence lower bound in order to understand how a VAE works. I recommend this tutorial.