What can we learn from this visualization?

Question

In here, I learned that the learned MNIST manifold can be visualized as the image below (on page 10, figure 4(b)).

My understanding for this visualization is:

We start from probability integral transform which allows us to sample from a continuous distribution by sampling from a uniform distribution on $[0, 1]^d$.

In this example, we want to sample from the prior $p_{\theta}(z)$ which is a bivariate normal distribution, so we can sample from a uniform distribution on a unit square $[0, 1] \times [0, 1] $, and then transform the sampled points to $z$ by the inverse CDF of $p_{\theta}(z)$.

So, in the picture, the square is actually a unit square, and those points that the digits sit are the points that we sampled from this uniform distribution. Since each point can be converted to a $z$, which is a point in the manifold $p_{\theta}(x)$. And each such $z$ can be further decoded by the trained $p_{\theta}(x|z)$. So, we have those reconstructed digits sitting on those points (those digits are means of $p_{\theta}(x|z)$).

But I don't know whether this understanding is correct or not.

Moreover, what can we learn from this visualization?

Answer 1

Your understanding is very good, save for a couple minor details:

• to create the image we don't sample from a uniform distribution on $I=[0,1]×[0,1]$. We actually define a uniform $20\times20$ grid in $I$, and we compute the values of $\mathbf{z}=(z_1,z_2)$ using the inverse of the Gaussian CDF (note that $z_1$ and $z_2$ are assumed independent, zero-mean and unit-variance here). If you use the term sampling, one would expect the initial values to be drawn randomly from the bivariate uniform distribution on $I$, and we wouldn't have the nice-looking grid.

• the "learned data manifold" Kingma & Welling are talking of, is the the decoder $p_{\boldsymbol{\theta}}(\mathbf{x}\vert\mathbf{z})$. To be more precise, it's the two vector functions $\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{z}):\mathbb{R}^2\to\mathbb{R}^{784}$ and $\log{\boldsymbol{\sigma}}_{\boldsymbol{\theta}}^2(\mathbf{z}):\mathbb{R}^2\to\mathbb{R}^{784}$, whose expressions in the case of a Gaussian MLP are:

$$ \begin{align} \boldsymbol{\mu}(\mathbf{z}) & = \mathbf{W}_2(\tanh({\mathbf{W}_1}\mathbf{z} + \mathbf{b}_1))+\mathbf{b}_2 \\ \log{\boldsymbol{\sigma}^2}(\mathbf{z}) & = \mathbf{W}_3(\tanh({\mathbf{W}_1}\mathbf{z} + \mathbf{b}_1))+\mathbf{b}_3 \end{align} $$

However, as you correctly realized, the visualization most likely shows only the 400 values (corresponding to the grid of 400 $\mathbf{z}$ values) of the mean function $\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{z})$. This is a dirty little VAE trick - when creating visualizations, people often don't really sample from $p_{\boldsymbol{\theta}}(\mathbf{x}\vert\mathbf{z}_i)$, but they just show the means of the distributions, to get better-looking images ¹. Since $\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{z})$ is the output of a neural network, it's a smooth function of $\mathbf{z}$, thus you get the nice-looking, smooth visualization in Figure 4.

What can we learn from this visualization?

We can see what features of the input space (the digits) have been coded into the latent variables, i.e., we can visualize the representation of the input space in terms of latent variables. Here the latent space has dimension 2: in a ideal world, we would like $z_1$ to encode some specific property of the images (for example, which digit we are generating, 0, 1, 2, etc.) and $z_2$ to encode a separate, independent property (for example, the slanting angle with which the digit has been drawn). This would be called a disentangled representation, and it's of interest because if we have one, then we can very easily generate images with some specific properties (for example, woman with long/short hair, man with/without sunglasses, etc.), by "fine-tuning" one or more of the latent variables.

If you use a vanilla VAE, it's very seldom the case that you learn a disentangled representation (you'd typically need $\beta$-VAE or InfoVAE/WassersteinVAE). In most of the cases, each latent variable affects all of the visual properties of the generated images, making it very difficult to perform the "fine-tuning" mentioned before. For example, in our case you can clearly see that neither $z_1$ or $z_2$ alone determine which digit is drawn. If you fix one variable and change only the other one, you can't generate all of the digits: to do so, you must vary both variables at the same time. The degree of slanting, instead, seems to be (very roughly) controlled by $z_2$: values of $z_2$ close to the center of the square seem to lead to more vertical digits, while values closer to the extremes seem to lead to more slanted digits, though the level of slanting is also affected (to a minor degree) by the current value of $z_1$.

Finally, one last (unsurprising) thing that we can learn, is that in the latent space two digits are close together if they're visually similar, not if they're consecutive in base 10. So 3 is close to 8, and also to 2 and 5, but it's not close to 4, 1 or 6, for example. As I said, this was to be expected, but still, it's nice to have a visual confirmation that our VAE trained well.

---

¹_{Of course if you explicitly say that you're sampling from your generative model, then in that case you're really drawing a sample $\mathbf{x}\sim p_{\boldsymbol{\theta}}(\mathbf{x}\vert\mathbf{z}_i)$, i.e., you're using $\log{\boldsymbol{\sigma}}_{\boldsymbol{\theta}}^2(\mathbf{z}_i)$ too (see for example Figure 5 of the same paper).}