Variational Autoencoder

A variational autoencoder (VAE) provides a probabilistic manner for describing an observation in latent space. Thus, rather than building an encoder that outputs a single value to describe each latent state attribute, we'll formulate our encoder to describe a probability distribution for each latent attribute. [Source]

Motivation, Objective and Related works

Motivation

Standard and variational autoencoders learn to represent the input just in a compressed form called the latent space or the bottleneck.
Therefore, the latent space formed after training the model is not necessarily continuous and, in effect, might not be easy to interpolate.
While these attributes explain the image and can be used in reconstructing the image from the compressed latent space, they do not allow the latent attributes to be expressed in a probabilistic fashion.

To provide an example, let's suppose we've trained an autoencoder model on a large dataset of faces with a encoding dimension of 6.
An ideal autoencoder will learn descriptive attributes of faces such as skin color, whether or not the person is wearing glasses, etc. in an attempt to describe an observation in some compressed representation.

Objectives

Variational autoencoders deal with this specific topic and express their latent attributes as a probability distribution, leading to the formation of a continuous latent space that can be easily sampled and interpolated.
When fed the same input, a variational autoencoder would construct latent attributes.
The latent attributes are then sampled from the latent distribution formed and fed to the decoder, reconstructing the input.
The motivation behind expressing the latent attributes as a probability distribution can be very easily understood via statistical expressions.

We may prefer to represent each latent attribute as a range of possible values.
Using a variational autoencoder, we can describe latent attributes in probabilistic terms.
When decoding from the latent state, we'll randomly sample from each latent state distribution to generate a vector as input for our decoder model.

Model

By constructing our encoder model to output a range of possible values (a statistical distribution) from which we'll randomly sample to feed into our decoder model, we're essentially enforcing a continuous, smooth latent space representation.
For any sampling of the latent distributions, we're expecting our decoder model to be able to accurately reconstruct the input. Thus, values which are nearby to one another in latent space should correspond with very similar reconstructions.

For variational autoencoder, the encoder model is sometimes referred to as the recognition model whereas the decoder model is sometimes referred to as the generative model.

We aim at identifying the characteristics of the latent vector z that reconstructs the output given a particular input. Effectively, we want to study the characteristics of the latent vector given a certain output x[p(z|x)].
While estimating the distribution becomes impossible mathematically, a much simpler and easier option is to build a parameterized model that can estimate the distribution. It does by minimizing the KL divergence between the original distribution and our parameterized one.
Expressing the parameterized distribution as q, we can infer the possible latent attributes used in the image reconstruction.
Assuming the prior z to be a multivariate Gaussian model, we can build a parameterized distribution as one containing two parameters, the mean and the variance. The corresponding distribution is then sampled and fed to the decoder, which then proceeds to reconstruct the input from the sample points.

Code

class VAE(torch.nn.Module):

def __init__(self, latent_dim=20):

super(VAE, self).__init__()

self.encoder = Encoder(latent_dim)

self.decoder = Decoder(latent_dim)

def forward(self,inputs):

z_mean, z_log_var = self.encoder(inputs)

z = self.reparameterize(z_mean, z_log_var)

reconstructed = self.decoder(z)

return reconstructed, z_mean, z_log_var

def reparameterize(self, mu, log_var):

std = torch.exp(0.5 * log_var)

eps = torch.randn_like(std)

return mu + (eps * std)

Loss Functions

Reconstruction Loss

- Distance Loss.
- Measure the loss of information.

Regularization Loss

- Kulback-Leibler divergence
- Manage the regularization of hidden space.

def final_loss(reconstruction, train_x, mu, logvar):

BCE = torch.nn.BCEWithLogitsLoss(reduction='sum')(reconstruction, train_x)

KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

return BCE + KLD

Reparameterization trick

- Randomly sample a value ε from a unit Gaussian and then scale this by the latent distribution variance σ and shift it by the mean μ of the same.
- We have left behind the sampling process as something done outside what the back-propagation pipeline handles, and the sampled value ε acts just like another input to the model, that is fed at the bottleneck.