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?