What is an appropriate size for a latent space of (variational) autoencoders and how it varies with the features of the images?

Question

I am training an autoencoder and a variational autoencoder using satellite and streetview images. I have tested my program on standard datasets such as MNIST and CelebA. It seems that the latent space dimension needed for those applications are fairly small. For example, MNIST is 28x28x1 and CelebA is 64x64x3 and for both a latent space bottleneck of 50 would be sufficient to observe reasonably reconstructed image.

But for the autoencoder I am constructing, I needed a dimension of ~20000 in order to see features. The variational autoencoder is not working, and I only see a few blobs of fuzzy color.

For those who have experience with training the autoencoders with your own images, what could be the problem? Are certain features easier to compress than others? or does it look like a sample size problem? (i have 20000 images in training) Is there any rule of thumb for the the factor of compression? Thanks!

Here is a best example of what I have got with my VAE. I am using a ResNeXt architecture and the image dimension is 64x64x3, the latent space dimension is very large (18432).

Answer 1

You are asking about several things here and while related, solving one, will not necessarily "solve" your problem. Let's look at them separately:

1. Optimal dimension of the latent space.
2. Blurry reconstructions.
3. Optimal sample size.

Optimal dimension of the latent space.

I'm unaware of a one-fit-all way to find the optimal dimensionality of $z$ but an easy way is to try with different values and look at the likelihood on the test-set $log(p)$ - pick the lowest dimensionality that maximises it. This is a solution in tune with Deep Learning spirit :)

Second solution, maybe a little more grounded, is to decompose your training data with SVD and look at the spectrum of singular values. The number of non trivial (=above some small threshold) values will give you a rough idea of the number of latent dimensions you are going to need.

Finally, you could allow for a lot of z-dimensions but augment your loss function in such way, that the encoder will be forced to only use what it needs. This is sometimes called Sparsity promoting, L1 or Lasso-type regularisation and is also something that can help with overfitting. Take a look at arXiv:1812.07238.

Blurry reconstructions

This is a notorious problem with VAE's and while there are a lot of theories on why this happens, my take is that the reason is two fold.

First, the loss function. With typical Cross-Entropy or MSE losses, we have this blunt bottom of the function where the minimum is residing allowing for a lot of similar "good solutions". See arXiv:1511.05440 and especially https://openreview.net/forum?id=rkglvsC9Ym for an easy fix that seems to improve the quality/sharpness of the reconstructions.

Second, the blurriness comes from the Variational formulation itself. Handwavy explained, we are trying to model some very-very complex data (images in your case), with a "simple" isotropic Gaussian. Not surprisingly, the result will be something the best the model can do given this constraint. Recall that the loss function in VAE's is called ELBO - Evidence Lower Bound - which basically tells us that we are trying to model a Lower Bound as best as we can and not the "actual data" distribution. Typically, introduction of a KL-multiplier, $\beta$, which relaxes the influence of the Gaussian prior, will give you better reconstructions (see $\beta$-VAE's).

Finally, if you are feeling especially adventurous, take a look at discrete VAE's (VQ-VAE's), which seem to have reconstructions on pair with GAN's. Sampling from them is not trivial, however.

Optimal sample size

As for optimal sample size, just choose an architecture that will not overfit. Decrease the number of neurons/layers, check your $\log(p)$ on the test-set, introduce Dropout, all the usual stuff.

Answer 2

From my experience (on MNIST digits), even when using a latent space of only $10$ nodes, the decoded reconstruction was pretty much ok. perhaps the architecture itself lacks the capabilities of encoding/decoding properly. Check out this summary and see if you can improve your results using a similar approach.