The generative story describes how each image sample is generated. The story is as follows - (a) Sample z ~ N(z | 0, I); (b) Sample x ~ N(x | f_mu(z), f_sig(z))
For any generative story, going forward tells us about test time, and going backwards helps us learn the parameters (note that last step has x
, our data).
Every model comes with its assumptions. The implicit assumption in VAE is : z-space is cleaved in such a way that different regions give different digits. Our goal is to figure out these regions. This is the "Encoder part".
We need to optimize the log-marginal-likelihood $logP(x)=log \Pi_i\int_z p(x_i | z)p(z)dz $. Let's say x_i
is 7
. Most of $p(x_i | z)$'s are going to be zero -- because most of the z
's don't even give 7
. We have "assumed" that only certain regions of z-space give 7
. You are cluelessly computing $p(x_i | z)$ for z
's that belong to 1,2,3
etc clusters (because of $\int_z$ part). Encoder (parameterized by $\phi$) helps us to not be clueless. At training time, you encode x_i
to z-space, and then decode it back to x-space. The math looks like this --
$$logP(x)=log \Pi_i\int_z \frac{p(x_i | z)q_{\phi}(z | x)p(z)}{q_{\phi}(z | x)}dz = \Sigma_i log\int_z \frac{p(x_i | z)q_{\phi}(z | x)p(z)}{q_{\phi}(z | x)}dz
$$
Notice that we have an expectation wrt encoder q_phi
inside log. So use Jensen inequality log(Expectation) >= Expectation(log) to get
$$log P(x) \geq \Sigma_i E_{q_{\phi}}[log\frac{p(x_i, z)}{q_{\phi}(z|x)} ]$$
where RHS is the ELBO term. Effectively, (x_i)----$q_{\phi}$----(z)----$p_{\theta}$----(x_i)
Intuitively, the encoder q
figures out how to take x
to z
in such a way that z
is "rich" (cleaved), and the decoder p
(parameterized by $\theta$) figures out to take z
from this "rich" z-space to x
.
At test time, you can't use encoder to take x
to z
. You don't even have x
- your target is to generate x
. So, you sample some random z
from N(z | 0, I). From your question, let's say $z = [0.1, 0.5]$. The decoder $p_{\theta}$ "knows" that this z-space is rich, and the z
you've sampled belongs to the 3
cluster (say). Hence, we generate the digit x = 3
.
You sample some other z
. The decoder figures that this belongs to the 1
cluster, and it generates x = 1
.
Can just a 2-dimensional z-space be "rich" enough so that different regions correspond to each of the 10 digits? Maybe not. You trained your model assuming that it can.