VAE implementation

We already have the implementation of a simple autoencoder.

To create a Variational Autoencoder (VAE), we need to modify the autoencoder's loss function by adding a regularization term. This term ensures that the distribution of the whole latent space values follows a standard Gaussian distribution with zero mean and unit variance.

Note

We will use KL divergence loss to encourage the latent distribution to approximate a Gaussian distribution.

KL divergence, or Kullback-Leibler divergence, is a measure from information theory that quantifies the difference between two probability distributions. It represents the amount of information lost when one probability distribution is used to approximate another.

Additionally, instead of generating the latent space values directly, our encoder will produce the parameters of a Gaussian distribution: the mean and variance vectors. These parameters will then be used to generate Gaussian samples, forming a probabilistic latent space.

Note

You can find the source code via the following Link. If you want to run the code or even change some components, you can copy the notebook and work with the copy.

Finally, the generated images are clear and understandable to the human eye. However, we still have one major problem: we can't control which particular digit is generated by the network. We simply pass some random sample to the decoder, and it generates a digit from 0 to 9.

If we want to generate a specific digit, we need to extend our architecture by using a Conditional Variational Autoencoder (CVAE).